TAKE: Trajectory-Aware Knowledge Estimation for Text Dataset Distillation

# Trajectory-Aware Knowledge Estimation for Text Dataset Distillation
Source: [https://arxiv.org/html/2607.11898](https://arxiv.org/html/2607.11898)
\(March 2026\)

###### Abstract

Large\-scale text corpora have become a quiet bottleneck in modern NLP, not just in storage, but in the accumulated cost of training, fine\-tuning, and continual learning\. We propose a text dataset distillation framework that reduces corpora to as little as 0\.1% of their original size while preserving downstream task fidelity\. We approach distillation through the lens of influence functions, which quantify each sample’s contribution to the downstream objective, a natural and principled basis for selection\. We introduceTrajectory\-Aware Knowledge Estimation\(TAKE\), which convolves the knowledge\-based influence along the training trajectory into a single per\-sample knowledge score, capturing informative samples\. These scores serve as sample weights within a discrete Optimal Transport objective, guiding prototype selection from a synthetically generated candidate pool\. We evaluate TAKE on downstream accuracy across text classification and natural language inference tasks at extreme compression \(0\.1% or 20 samples/class\), showing that data efficiency is achievable without sacrificing task fidelity\. The approach is theoretically grounded, with broader implications for coreset construction and data\-centric AI\. We release our source code at[https://github\.com/votrinhan88/take](https://github.com/votrinhan88/take)\.

Keywords:Text dataset distillation⋅\\cdotInfluence functions⋅\\cdotOptimal transport\.

## 1Introduction

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](https://arxiv.org/html/2607.11898#bib.bib1)\], 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](https://arxiv.org/html/2607.11898#bib.bib1),[2](https://arxiv.org/html/2607.11898#bib.bib2),[3](https://arxiv.org/html/2607.11898#bib.bib3)\]\. 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](https://arxiv.org/html/2607.11898#bib.bib4),[5](https://arxiv.org/html/2607.11898#bib.bib5),[6](https://arxiv.org/html/2607.11898#bib.bib6),[7](https://arxiv.org/html/2607.11898#bib.bib7),[8](https://arxiv.org/html/2607.11898#bib.bib8)\]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](https://arxiv.org/html/2607.11898#S2.SS1)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 proposeTrajectory\-Aware Knowledge Estimation\(TAKE\), the first text DD method to satisfy all three desiderata\. We address the first via influence functions\[[9](https://arxiv.org/html/2607.11898#bib.bib9)\], which quantify each sample’s counterfactual effect on the downstream loss and provide a theoretically grounded, task\-aligned weighting scheme \(§[3](https://arxiv.org/html/2607.11898#S3)\)\. However, scoring at a single checkpoint introduces thehard\-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](https://arxiv.org/html/2607.11898#S2)reviews related work; §[3](https://arxiv.org/html/2607.11898#S3)presents theoretical motivation; §[4](https://arxiv.org/html/2607.11898#S4)details TAKE and our DD pipeline; §[5](https://arxiv.org/html/2607.11898#S5)reports experiments and discussion; §[6](https://arxiv.org/html/2607.11898#S6)reflects limitations and ethical concerns; §[7](https://arxiv.org/html/2607.11898#S7)concludes\.

## 2Related Works

### 2\.1Text Dataset Distillation

Dataset distillation was introduced by\[[1](https://arxiv.org/html/2607.11898#bib.bib1)\]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](https://arxiv.org/html/2607.11898#bib.bib2)\], distribution matching\[[10](https://arxiv.org/html/2607.11898#bib.bib10)\], and trajectory matching\[[3](https://arxiv.org/html/2607.11898#bib.bib3),[11](https://arxiv.org/html/2607.11898#bib.bib11)\]\. 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](https://arxiv.org/html/2607.11898#bib.bib4),[5](https://arxiv.org/html/2607.11898#bib.bib5)\]directly adapted the meta\-learning formulation\[[1](https://arxiv.org/html/2607.11898#bib.bib1)\], 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](https://arxiv.org/html/2607.11898#bib.bib6)\]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](https://arxiv.org/html/2607.11898#bib.bib7)\]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](https://arxiv.org/html/2607.11898#bib.bib8)\]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](https://arxiv.org/html/2607.11898#S2.T1)maps this progression: each generation addressed some gaps while leaving others open\.

Table 1:Landscape of text DD methods\.GapSLDDDDTCDDALDiLMDaLLMETAKETask alignment✓\\checkmark✓\\checkmark△\\triangle△\\triangle−\-✓\\checkmarkGlobal optimization△\\triangle△\\triangle△\\triangle△\\triangle△\\triangle✓\\checkmarkInterpretability△\\triangle−\-−\-✓\\checkmark✓\\checkmark✓\\checkmarkImportance weighting−\-−\-−\-−\-−\-✓\\checkmark−\-: absent,△\\triangle: partial,✓\\checkmark: addressedA 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\.2Influence Functions and the Hard\-Sample Bias

Influence functions \(IF\)\[[12](https://arxiv.org/html/2607.11898#bib.bib12),[9](https://arxiv.org/html/2607.11898#bib.bib9)\]quantify the leave\-one\-out effect of a training point viaℐ​\(z\)=∇⊤Hθ−1​∇\\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](https://arxiv.org/html/2607.11898#bib.bib13),[14](https://arxiv.org/html/2607.11898#bib.bib14),[15](https://arxiv.org/html/2607.11898#bib.bib15)\]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](https://arxiv.org/html/2607.11898#bib.bib16)\]show that networks fit clean patterns before noisy ones; dataset cartography\[[17](https://arxiv.org/html/2607.11898#bib.bib17)\]confirms that single\-checkpoint statistics conflate difficulty with noise; and Co\-teaching\[[18](https://arxiv.org/html/2607.11898#bib.bib18)\]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\.3Optimal Transport for Distribution Matching

Distribution matching has been proposed as a tractable DD surrogate\[[19](https://arxiv.org/html/2607.11898#bib.bib19)\]: 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](https://arxiv.org/html/2607.11898#Thmtheorem1)\)\. Common divergences have practical shortcomings at the extreme budgets relevant here\. MMD\[[20](https://arxiv.org/html/2607.11898#bib.bib20)\]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](https://arxiv.org/html/2607.11898#bib.bib21)\]\.

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](https://arxiv.org/html/2607.11898#bib.bib22)\]makes this tractable at corpus scale\. OT has been widely applied in NLP and machine learning\[[23](https://arxiv.org/html/2607.11898#bib.bib23),[24](https://arxiv.org/html/2607.11898#bib.bib24),[25](https://arxiv.org/html/2607.11898#bib.bib25),[26](https://arxiv.org/html/2607.11898#bib.bib26)\], but never to text DD, and never combined with a knowledge\-reweighted source distribution in any DD setting\. TAKE fills both gaps:κn\\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\.

## 3Theoretical Motivation

### 3\.1Reweighted Distribution Matching for Dataset Distillation

The dataset distillation \(DD\) objective seeks𝒟~\\tilde\{\\mathcal\{D\}\}with\|𝒟~\|=K≪N\|\\tilde\{\\mathcal\{D\}\}\|=K\\ll Nminimizing the loss on the original data:

min𝒟~⁡𝔼θ0∼p​\(θ0\)​\[𝔼z∼𝒟​\[ℓ​\(z;F​\(θ0;𝒟~,η,T\)\)\]\]\\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\]\(1\)whereF​\(θ0;𝒟~,η,T\)F\(\\theta\_\{0\};\\tilde\{\\mathcal\{D\}\},\\eta,T\)denotes the model obtained by runningTTsteps of gradient descent with learning rateη\\etaon𝒟~\\tilde\{\\mathcal\{D\}\}from initializationθ0\\theta\_\{0\}\. This bi\-level problem is intractable for discrete text\. We relax it to distribution matching via the gradient\-field equivalence argument\[[19](https://arxiv.org/html/2607.11898#bib.bib19)\]: if𝒟~\\tilde\{\\mathcal\{D\}\}induces the same expected gradient field as𝒟\\mathcal\{D\}throughout training, the two runs converge to comparable parameters\. Since∇θ𝔼z∼P​\[ℓ​\(z;θ\)\]=𝔼z∼P​\[∇θℓ​\(z;θ\)\]\\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](https://arxiv.org/html/2607.11898#S3.E1)\. This yields the distribution matching \(DM\) objective:

min\|𝒟~\|=K⁡d​\(P𝒟~,P𝒟\)\\min\_\{\|\\tilde\{\\mathcal\{D\}\}\|=K\}\\;d\\\!\\left\(P\_\{\\tilde\{\\mathcal\{D\}\}\},\\;P\_\{\\mathcal\{D\}\}\\right\)\(2\)for a divergenceddin a task\-relevant feature space\. Eq\.[2](https://arxiv.org/html/2607.11898#S3.E2)is a surrogate, not an equivalence; the formal gap bound appears as Theorem[1](https://arxiv.org/html/2607.11898#Thmtheorem1)onceddis specified\.

At extreme compression \(ρ=K/N≪1\\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: definewn≥0w\_\{n\}\\geq 0,∑nwn=1\\sum\_\{n\}w\_\{n\}=1, inducingPw=∑nwn​δznP^\{w\}=\\sum\_\{n\}w\_\{n\}\\delta\_\{z\_\{n\}\}, yielding the reweighted distribution matching \(RDM\) objective:

min\|𝒟~\|=K⁡d​\(P𝒟~,Pw\)\\min\_\{\|\\tilde\{\\mathcal\{D\}\}\|=K\}\\;d\\\!\\left\(P\_\{\\tilde\{\\mathcal\{D\}\}\},\\;P^\{w\}\\right\)\(3\)The weightswnw\_\{n\}must reflect downstream learning contribution, not geometry alone, motivating gradient\-based estimation\. Influence self\-scores \(§[2\.2](https://arxiv.org/html/2607.11898#S2.SS2)\) are the natural candidate: they are task\-aligned by construction and theoretically grounded\. We now show that the simplest instantiation—wn∝ℐself​\(zn\)w\_\{n\}\\propto\\mathcal\{I\}\_\{\\text\{self\}\}\(z\_\{n\}\)at a single checkpoint, select top\-KK—fails on both counts: the weights are biased toward noisy samples \(Proposition[1](https://arxiv.org/html/2607.11898#Thmproposition1)\), and independent top\-KKselection ignores distributional coverage \(Remark[3](https://arxiv.org/html/2607.11898#Thmremark3)\)\.

### 3\.2The Naive Instantiation Fails: Two Open Problems

Define the influence self\-score at checkpointθ^\\hat\{\\theta\}as

ℐself​\(zn;θ^\)=∇θℓ​\(zn;θ^\)⊤​Hθ^−1​∇θℓ​\(zn;θ^\),\\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\}\),\(5\)whereHθ^=1N​∑n∇θ2ℓ​\(zn;θ^\)H\_\{\\hat\{\\theta\}\}=\\frac\{1\}\{N\}\\sum\_\{n\}\\nabla^\{2\}\_\{\\theta\}\\ell\(z\_\{n\};\\hat\{\\theta\}\)is the empirical Hessian, specializing\[[9](https://arxiv.org/html/2607.11898#bib.bib9)\]by setting the test point equal to the training point\. Highℐself\\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: setwn∝ℐself​\(zn;θ^\)w\_\{n\}\\propto\\mathcal\{I\}\_\{\\text\{self\}\}\(z\_\{n\};\\hat\{\\theta\}\)and returnarg⁡max\|𝒜\|=K​∑i∈𝒜wi\\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\.

###### Proposition 1\(Hard\-Sample Bias\)\.

Let𝒟clean,𝒟noisy\\mathcal\{D\}\_\{\\text\{clean\}\},\\,\\mathcal\{D\}\_\{\\text\{noisy\}\}partition𝒟\\mathcal\{D\}, under the assumptions:*\(i\)*‖∇θℓ​\(z;θ^\)‖≤ϵ\\\|\\nabla\_\{\\theta\}\\ell\(z;\\hat\{\\theta\}\)\\\|\\leq\\epsilonforz∈𝒟cleanz\\in\\mathcal\{D\}\_\{\\text\{clean\}\}, i\.e\., clean samples have near\-zero gradient at convergence;*\(ii\)*‖∇θℓ​\(z;θ^\)‖≥δ≫ϵ\\\|\\nabla\_\{\\theta\}\\ell\(z;\\hat\{\\theta\}\)\\\|\\geq\\delta\\gg\\epsilonforz∈𝒟noisyz\\in\\mathcal\{D\}\_\{\\text\{noisy\}\}, i\.e\., noisy samples retain large residual gradients\. Letcond​\(H\)=λmax​\(H\)/λmin​\(H\)\\mathrm\{cond\}\(H\)=\\lambda\_\{\\max\}\(H\)/\\lambda\_\{\\min\}\(H\)denotes the condition number of the HessianHH\. Then:

ΔHSB:=𝔼z∼𝒟noisy​\[ℐself​\(z;θ^\)\]−𝔼z∼𝒟clean​\[ℐself​\(z;θ^\)\]≥δ2−cond​\(Hθ^\)​ϵ2λmax​\(Hθ^\)\>0\\Delta\_\{\\textup\{HSB\}\}:=\\mathbb\{E\}\_\{z\\sim\\mathcal\{D\}\_\{\\textup\{noisy\}\}\}\[\\mathcal\{I\}\_\{\\textup\{self\}\}\(z;\\hat\{\\theta\}\)\]\-\\mathbb\{E\}\_\{z\\sim\\mathcal\{D\}\_\{\\textup\{clean\}\}\}\[\\mathcal\{I\}\_\{\\textup\{self\}\}\(z;\\hat\{\\theta\}\)\]\\;\\geq\\;\\frac\{\\delta^\{2\}\-\\mathrm\{cond\}\(H\_\{\\hat\{\\theta\}\}\)\\,\\epsilon^\{2\}\}\{\\lambda\_\{\\max\}\(H\_\{\\hat\{\\theta\}\}\)\}\\;\>\\;0\(6\)

###### Proof sketch \(full proof: Appendix A\.1\)\.

Sinceℐself​\(z\)=∇⊤Hθ^−1​∇\\mathcal\{I\}\_\{\\text\{self\}\}\(z\)=\\nabla^\{\\top\}H\_\{\\hat\{\\theta\}\}^\{\-1\}\\nabla, spectral bounds onHθ^−1H\_\{\\hat\{\\theta\}\}^\{\-1\}give‖∇‖2/λmax​\(Hθ^\)≤ℐself​\(z\)≤‖∇‖2/λmin​\(Hθ^\)\\\|\\nabla\\\|^\{2\}/\\lambda\_\{\\max\}\(H\_\{\\hat\{\\theta\}\}\)\\leq\\mathcal\{I\}\_\{\\text\{self\}\}\(z\)\\leq\\\|\\nabla\\\|^\{2\}/\\lambda\_\{\\min\}\(H\_\{\\hat\{\\theta\}\}\)\. Applying the lower bound to noisy samples and the upper bound to clean samples yields Eq\.[6](https://arxiv.org/html/2607.11898#S3.E6)\. Positivity requiresδ/ϵ\>cond​\(Hθ^\)\\delta/\\epsilon\>\\sqrt\{\\mathrm\{cond\}\(H\_\{\\hat\{\\theta\}\}\)\}, which is mild givenδ≫ϵ\\delta\\gg\\epsilon\. ∎

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\-KKselection is blind to distributional coverage\.

The two failures compound: biasedwnw\_\{n\}misdirects a population\-level solver toward noisy regions; correctwnw\_\{n\}fed into an independent selector still produces redundancy\. They respectively yield two open problems:

1. \(a\)Biased weighting\.Estimatewnw\_\{n\}by integrating gradient signals across trajectory\{θt\}t=0T−1\\\{\\theta\_\{t\}\\\}\_\{t=0\}^\{T\-1\}, reducingΔHSB\\Delta\_\{\\textup\{HSB\}\}—which we will address in §[4\.2](https://arxiv.org/html/2607.11898#S4.SS2)\(formalized as Proposition[2](https://arxiv.org/html/2607.11898#Thmproposition2); Appendix A\.2\)\.
2. \(b\)Independent selection\.Solve Eq\.[3](https://arxiv.org/html/2607.11898#S3.E3)directly for a metricddthat jointly optimises coverage ofPwP^\{w\}—which we will address in §[4\.4](https://arxiv.org/html/2607.11898#S4.SS4)\.

## 4Methodology

### 4\.1Problem Setup and Framework Overview

TAKE instantiates the RDM objective \(Eq\.[3](https://arxiv.org/html/2607.11898#S3.E3)\) by solving the two open problems from §[3\.2](https://arxiv.org/html/2607.11898#S3.SS2)\. Let𝒟=\{zn\}n=1N\\mathcal\{D\}=\\\{z\_\{n\}\\\}\_\{n=1\}^\{N\}denote the training corpus,𝒟synth=\{z~c\}c=1M\\mathcal\{D\}\_\{\\text\{synth\}\}=\\\{\\tilde\{z\}\_\{c\}\\\}\_\{c=1\}^\{M\}the synthetic candidate pool, and𝒟~⊂𝒟synth\\tilde\{\\mathcal\{D\}\}\\subset\\mathcal\{D\}\_\{\\text\{synth\}\}the distilled set of budgetKK\. Indicesn∈\[N\]n\\in\[N\]andt∈\[T\]t\\in\[T\]range over training samples and checkpoints respectively\. TAKE proceeds in three stages:

1. 1\.Score—train a logistic probe on𝒟\\mathcal\{D\}acrossTTcheckpoints, compute influence matrixI∈ℝN×TI\\in\\mathbb\{R\}^\{N\\times T\}, apply reciprocal reweighting and temporal weighting to obtain trajectory\-aware knowledge scoresκn\\kappa\_\{n\}\.
2. 2\.Generate—fine\-tune a small language model on𝒟\\mathcal\{D\}and sample to produce the synthetic candidate pool𝒟synth\\mathcal\{D\}\_\{\\text\{synth\}\}\.
3. 3\.Distill—solve discrete OT with source weightsκn\\kappa\_\{n\}over𝒟synth\\mathcal\{D\}\_\{\\text\{synth\}\}to obtain the distilled corpus𝒟~\\tilde\{\\mathcal\{D\}\}of budgetKK\.

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

### 4\.2Gradient\-Based Influence Along the Trajectory

The hard\-sample bias \(HSB; §[3\.2](https://arxiv.org/html/2607.11898#S3.SS2)\) 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κn\\kappa\_\{n\}per sample\. We buildκn\\kappa\_\{n\}in two steps: constructing the influence matrixII\(§[4\.2\.1](https://arxiv.org/html/2607.11898#S4.SS2.SSS1)\), then collapsing it into a per\-sample scalar score via reciprocal reweighting and temporal weighting \(§[4\.2\.2](https://arxiv.org/html/2607.11898#S4.SS2.SSS2)\)\. The resultingκn\\kappa\_\{n\}provides a task\-aligned, bias\-corrected weight forPwP^\{w\}\.

#### 4\.2\.1Influence Matrix

Our goal is to constructI∈ℝN×TI\\in\\mathbb\{R\}^\{N\\times T\}whose entryI\(n,t\)I\_\{\(n,t\)\}reflects how much sampleznz\_\{n\}contributes to learning at checkpointθt\\theta\_\{t\}\. The influence self\-score \(§[3\.2](https://arxiv.org/html/2607.11898#S3.SS2), Eq\.[5](https://arxiv.org/html/2607.11898#S3.E5)\) 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](https://arxiv.org/html/2607.11898#bib.bib13)\]to the self\-influence setting, querying training points instead of test points:

I\(n,t\)=‖∇θℓ​\(zn;θt\)‖\.I\_\{\(n,t\)\}=\\\|\\nabla\_\{\\theta\}\\ell\(z\_\{n\};\\,\\theta\_\{t\}\)\\\|\.\(7\)As theH=IH=Ispecial case of the influence self\-score \(Remark[2](https://arxiv.org/html/2607.11898#Thmremark2)\),I\(n,t\)I\_\{\(n,t\)\}exhibits the same HSB structure—the bias thatκn\\kappa\_\{n\}corrects via reciprocal inversion and trajectory integration \(Proposition[2](https://arxiv.org/html/2607.11898#Thmproposition2)\)\. Scores are column\-wise normalizedI\(n,t\)←I\(n,t\)/∑n′I\(n′,t\)I\_\{\(n,t\)\}\\leftarrow I\_\{\(n,t\)\}/\\sum\_\{n^\{\\prime\}\}I\_\{\(n^\{\\prime\},t\)\}for comparability across checkpoints\.

In practice, we score samples with a linear headWWatop a pre\-trained, frozen sentence encoderϕ\\phi, givingθ=\(W,ϕ\)\\theta=\(W,\\phi\)\. The signal quality comes not from the linear head but fromϕ​\(xn\)\\phi\(x\_\{n\}\):ϕ\\phiis pre\-trained on large\-scale corpora andWWis 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](https://arxiv.org/html/2607.11898#bib.bib27)\], we approximate full\-model influence via the last layer—here, the linear headWW:

∇θℓ​\(zn;θt\)≈∇Wℓ​\(zn;θt\),\\nabla\_\{\\theta\}\\ell\(z\_\{n\};\\,\\theta\_\{t\}\)\\approx\\nabla\_\{W\}\\ell\(z\_\{n\};\\,\\theta\_\{t\}\),\(8\)trading a small bias for tractability\. By the Goodfellow outer\-product identity\[[28](https://arxiv.org/html/2607.11898#bib.bib28)\],∇Wℓ=gn​ϕ​\(xn\)⊤\\nabla\_\{W\}\\ell=g\_\{n\}\\phi\(x\_\{n\}\)^\{\\top\}wheregng\_\{n\}is the output gradient\. The Goodfellow trick vectorizes the batch backward pass, avoiding per\-sample loops\. Adapting to a new task is trivial: retrainWWon the new corpus;ϕ\\phiand the influence pipeline are unchanged\.

#### 4\.2\.2Knowledge 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 whatI\(n,t\)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 kernelk:\{0,…,T−1\}→ℝ\>0k:\\\{0,\\ldots,T\-1\\\}\\to\\mathbb\{R\}\_\{\>0\}, giving:

κn=∑t=0T−1k​\(t\)I\(n,t\),\\kappa\_\{n\}=\\sum\_\{t=0\}^\{T\-1\}\\frac\{k\(t\)\}\{I\_\{\(n,t\)\}\},\(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](https://arxiv.org/html/2607.11898#bib.bib16)\]\. Thus, concentrating mass on the early\-to\-mid phase and discounting late memorization directly corrects the HSB structure identified in Proposition[1](https://arxiv.org/html/2607.11898#Thmproposition1)\. By default we use the exponential kernelk​\(t\)=e−λ​t/Tk\(t\)=e^\{\-\\lambda t/T\}\(λ\>0\\lambda\>0\), withλ\\lambdatreated as a hyperparameter\. Alternative kernels \(linear, cosine\) are drop\-in replacements, and Proposition[2](https://arxiv.org/html/2607.11898#Thmproposition2)guarantees bias reduction holds for any strictly decreasingkk\. The resultingκn\\kappa\_\{n\}summarizes each sample’s cumulative, bias\-corrected learning contribution across the full trajectory, resolving open problem[\(a\)](https://arxiv.org/html/2607.11898#S3.I1.i1)\. Normalized source weightswn=κn/∑n′κn′w\_\{n\}=\\kappa\_\{n\}/\\sum\_\{n^\{\\prime\}\}\\kappa\_\{n^\{\\prime\}\}define the reweighted distributionPwP^\{w\}and serve as input to the OT selection step \(§[4\.4](https://arxiv.org/html/2607.11898#S4.SS4)\)\.

###### Proposition 2\(Trajectory Integration Reduces HSB\)\.

Under a smooth loss and bounded gradient noise,κn\\kappa\_\{n\}produces a strictly smaller noisy\-over\-clean gap than the single last\-checkpoint baselineκn†=1/I\(n,T−1\)\\kappa^\{\\dagger\}\_\{n\}=1/I\_\{\(n,T\-1\)\}:

𝔼z∼𝒟noisy​\[κn\]−𝔼z∼𝒟clean​\[κn\]<𝔼z∼𝒟noisy​\[κn†\]−𝔼z∼𝒟clean​\[κn†\]\.\\mathbb\{E\}\_\{z\\sim\\mathcal\{D\}\_\{\\textup\{noisy\}\}\}\[\\kappa\_\{n\}\]\-\\mathbb\{E\}\_\{z\\sim\\mathcal\{D\}\_\{\\textup\{clean\}\}\}\[\\kappa\_\{n\}\]<\\mathbb\{E\}\_\{z\\sim\\mathcal\{D\}\_\{\\textup\{noisy\}\}\}\[\\kappa^\{\\dagger\}\_\{n\}\]\-\\mathbb\{E\}\_\{z\\sim\\mathcal\{D\}\_\{\\textup\{clean\}\}\}\[\\kappa^\{\\dagger\}\_\{n\}\]\.\(10\)

###### Proof sketch \(full proof: Appendix A\.2\)\.

DefineΔt:=𝔼noisy​\[1/I\(n,t\)\]−𝔼clean​\[1/I\(n,t\)\]\\Delta\_\{t\}:=\\mathbb\{E\}\_\{\\textup\{noisy\}\}\[1/I\_\{\(n,t\)\}\]\-\\mathbb\{E\}\_\{\\textup\{clean\}\}\[1/I\_\{\(n,t\)\}\]as the noisy\-over\-clean gap\. The left\-hand side of Eq\.[10](https://arxiv.org/html/2607.11898#S4.E10)equals∑tk​\(t\)​Δt\\sum\_\{t\}k\(t\)\\Delta\_\{t\}\. Fort≤Tmt\\leq T\_\{m\}\(before memorization onset\): clean gradients decay soΔt≤0\\Delta\_\{t\}\\leq 0, strictly at somet∗t^\{\*\}\(Appendix A\.2\)\. BeyondTmT\_\{m\}: noisy gradients also drop \(memorization\), soΔT−1\>0\\Delta\_\{T\-1\}\>0— the last\-checkpoint baseline is inflated by memorization\. The decreasing kernel upweightst≤Tmt\\leq T\_\{m\}\(whereΔt<0\\Delta\_\{t\}<0\) and downweightst\>Tmt\>T\_\{m\}\(whereΔt\>0\\Delta\_\{t\}\>0\)\. Sincekkis strictly decreasing, the negative\-Δt\\Delta\_\{t\}early terms receive more weight and the large positiveΔT−1\\Delta\_\{T\-1\}term receives the least, making the weighted sum strictly smaller thanΔT−1\>0\\Delta\_\{T\-1\}\>0\. ∎

### 4\.3Synthetic Candidate Pool Generation

Rather than selecting directly from𝒟\\mathcal\{D\}, TAKE draws𝒟~\\tilde\{\\mathcal\{D\}\}from a synthetic candidate pool𝒟synth\\mathcal\{D\}\_\{\\text\{synth\}\}\. Direct subsampling has two problems: at extreme compression, verbatim records raise privacy concerns, and the budgetKKis too small to cover the full training support\. We generate𝒟synth\\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\.𝒟synth\\mathcal\{D\}\_\{\\text\{synth\}\}serves solely as the selection pool and is never appended to downstream training data\.

TAKE does*not*assumeP𝒟synth=P𝒟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 towardPwP^\{w\}regardless of pool distribution, soP𝒟~P\_\{\\tilde\{\\mathcal\{D\}\}\}approximatesPwP^\{w\}provided𝒟synth\\mathcal\{D\}\_\{\\text\{synth\}\}covers the support ofP𝒟P\_\{\\mathcal\{D\}\}\(Assumption[1](https://arxiv.org/html/2607.11898#Thmassumption1)\):

###### Assumption 1\(Pool Coverage\)\.

For every training samplezn∈𝒟z\_\{n\}\\in\\mathcal\{D\}, there exists a synthetic samplez~c∈𝒟synth\\tilde\{z\}\_\{c\}\\in\\mathcal\{D\}\_\{\\text\{synth\}\}within embedding distancerr:minc⁡‖ϕ​\(zn\)−ϕ​\(z~c\)‖≤r\\min\_\{c\}\\\|\\phi\(z\_\{n\}\)\-\\phi\(\\tilde\{z\}\_\{c\}\)\\\|\\leq r\.

Under Assumption[1](https://arxiv.org/html/2607.11898#Thmassumption1),W1​\(Pw,P𝒟synth\)≤rW\_\{1\}\(P^\{w\},P\_\{\\mathcal\{D\}\_\{\\text\{synth\}\}\}\)\\leq r, so the optimal transport cost is at mostrrbefore any prototype selection\. A practitioner can monitor coverage directly by reporting the mean nearest\-neighbor distancer¯=1N​∑nminc⁡‖ϕ​\(zn\)−ϕ​\(z~c\)‖\\bar\{r\}=\\frac\{1\}\{N\}\\sum\_\{n\}\\min\_\{c\}\\\|\\phi\(z\_\{n\}\)\-\\phi\(\\tilde\{z\}\_\{c\}\)\\\|\.

### 4\.4Distributional Prototype Selection via Discrete Optimal Transport

Both𝒟\\mathcal\{D\}and𝒟synth\\mathcal\{D\}\_\{\\text\{synth\}\}are finite sample sets, so the matching problem is naturally over empirical measures\. Discrete OT operates directly onμ=∑nwn​δzn\\mu=\\sum\_\{n\}w\_\{n\}\\delta\_\{z\_\{n\}\}andν=1\|𝒟synth\|​∑cδz~c\\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\.

LetCn​c=‖ϕ​\(zn\)−ϕ​\(z~c\)‖2C\_\{nc\}=\\\|\\phi\(z\_\{n\}\)\-\\phi\(\\tilde\{z\}\_\{c\}\)\\\|^\{2\}be the cost matrix in a task\-relevant embedding space\. Bare Kantorovich OT,minγ∈Π​\(μ,ν\)⁡⟨C,γ⟩\\min\_\{\\gamma\\in\\Pi\(\\mu,\\nu\)\}\\langle C,\\gamma\\rangle, produces a sparse, non\-differentiable plan and scales asO​\(N3\)O\(N^\{3\}\)\. We add entropic regularization to obtain a strictly convex, smooth objective solvable via Sinkhorn–Knopp iterations atO​\(N⋅\|𝒟synth\|\)O\(N\\cdot\|\\mathcal\{D\}\_\{\\text\{synth\}\}\|\)per step:

γ∗=arg⁡minγ∈Π​\(μ,ν\)⁡⟨C,γ⟩\+ε⋅KL​\(γ∥μ⊗ν\)\.\\gamma^\{\*\}=\\arg\\min\_\{\\gamma\\in\\Pi\(\\mu,\\nu\)\}\\langle C,\\gamma\\rangle\+\\varepsilon\\cdot\\mathrm\{KL\}\(\\gamma\\\|\\mu\\otimes\\nu\)\.\(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\}\}\. TheKKcandidates with highest received mass∑nγn​c∗\\sum\_\{n\}\\gamma^\{\*\}\_\{nc\}are selected as𝒟~\\tilde\{\\mathcal\{D\}\}, resolving open problem[\(b\)](https://arxiv.org/html/2607.11898#S3.I1.i2)\. The following theorem closes the theoretical chain from RDM \(Eq\.[3](https://arxiv.org/html/2607.11898#S3.E3)\) back to DD \(Eq\.[1](https://arxiv.org/html/2607.11898#S3.E1)\), bounding the downstream loss gap in terms of the OT cost\.

###### Theorem 1\(Distribution Matching Gap Bound\)\.

Letf𝒮∗f^\{\*\}\_\{\\mathcal\{S\}\}denote the model trained on𝒮\\mathcal\{S\}, and letℒ​\(𝒮\)=𝔼z∼Pw​\[ℓ​\(z;f𝒮∗\)\]\\mathcal\{L\}\(\\mathcal\{S\}\)=\\mathbb\{E\}\_\{z\\sim P^\{w\}\}\[\\ell\(z;\\,f^\{\*\}\_\{\\mathcal\{S\}\}\)\]\. Assume:*\(i\)*ℓ​\(⋅;f\)\\ell\(\\cdot;\\,f\)isLL\-Lipschitz in the embedding metric for allff;*\(ii\)*the model class has uniform stability constantβK=O​\(1/K\)\\beta\_\{K\}=O\(1/K\)\[[29](https://arxiv.org/html/2607.11898#bib.bib29)\], so thatsupz\|ℓ​\(z;f𝒟~∗\)−ℓ​\(z;f𝒟∗\)\|≤βK\\sup\_\{z\}\|\\ell\(z;\\,f^\{\*\}\_\{\\tilde\{\\mathcal\{D\}\}\}\)\-\\ell\(z;\\,f^\{\*\}\_\{\\mathcal\{D\}\}\)\|\\leq\\beta\_\{K\}\. Then:

\|ℒ​\(𝒟~\)−ℒ​\(𝒟\)\|≤2​L⋅W1​\(Pw,P𝒟~\)\+βK,\\bigl\|\\mathcal\{L\}\(\\tilde\{\\mathcal\{D\}\}\)\-\\mathcal\{L\}\(\\mathcal\{D\}\)\\bigr\|\\;\\leq\\;2L\\cdot W\_\{1\}\(P^\{w\},\\,P\_\{\\tilde\{\\mathcal\{D\}\}\}\)\+\\beta\_\{K\},\(12\)whereW1W\_\{1\}is the Wasserstein\-1 distance under the embedding metric\. AsK→∞K\\to\\infty,βK→0\\beta\_\{K\}\\to 0, and the bound is dominated by theW1W\_\{1\}term controlled by the OT objective\.

###### Proof sketch \(full proof: Appendix A\.3\)\.

Insert two intermediate terms,𝔼P𝒟~​\[ℓ​\(z;f𝒟~∗\)\]\\mathbb\{E\}\_\{P\_\{\\tilde\{\\mathcal\{D\}\}\}\}\[\\ell\(z;\\,f^\{\*\}\_\{\\tilde\{\\mathcal\{D\}\}\}\)\]and𝔼P𝒟~​\[ℓ​\(z;f𝒟∗\)\]\\mathbb\{E\}\_\{P\_\{\\tilde\{\\mathcal\{D\}\}\}\}\[\\ell\(z;\\,f^\{\*\}\_\{\\mathcal\{D\}\}\)\], and apply the triangle inequality twice\. This yields two distribution\-shift terms \(under fixedf𝒟~∗f^\{\*\}\_\{\\tilde\{\\mathcal\{D\}\}\}andf𝒟∗f^\{\*\}\_\{\\mathcal\{D\}\}respectively\), each bounded byL⋅W1​\(Pw,P𝒟~\)L\\cdot W\_\{1\}\(P^\{w\},P\_\{\\tilde\{\\mathcal\{D\}\}\}\)via Kantorovich–Rubinstein duality; and one model\-difference term underP𝒟~P\_\{\\tilde\{\\mathcal\{D\}\}\}, bounded byβK=O​\(1/K\)\\beta\_\{K\}=O\(1/K\)via uniform stability\. ∎

The bound is stated with respect toPwP^\{w\}rather than the uniformP𝒟P\_\{\\mathcal\{D\}\}— this is intentional, asPwP^\{w\}is the designed optimization target of the RDM objective; sincePwP^\{w\}reweightsP𝒟P\_\{\\mathcal\{D\}\}over the same support, minimizingW1​\(Pw,P𝒟~\)W\_\{1\}\(P^\{w\},P\_\{\\tilde\{\\mathcal\{D\}\}\}\)controls coverage ofP𝒟P\_\{\\mathcal\{D\}\}up to the reweighting gap bounded in Remark[1](https://arxiv.org/html/2607.11898#Thmremark1)\.

Together, Propositions[1](https://arxiv.org/html/2607.11898#Thmproposition1)–[2](https://arxiv.org/html/2607.11898#Thmproposition2)and Theorem[1](https://arxiv.org/html/2607.11898#Thmtheorem1)establish the full theoretical basis for TAKE: HSB is structural and unavoidable at a single checkpoint \(Proposition[1](https://arxiv.org/html/2607.11898#Thmproposition1)\); trajectory integration strictly reduces it \(Proposition[2](https://arxiv.org/html/2607.11898#Thmproposition2)\); and minimizing the OT cost directly controls the downstream loss gap underPwP^\{w\}\(Theorem[1](https://arxiv.org/html/2607.11898#Thmtheorem1); the gap toP𝒟P\_\{\\mathcal\{D\}\}is bounded by Remark[1](https://arxiv.org/html/2607.11898#Thmremark1)\)\. Algorithm[1](https://arxiv.org/html/2607.11898#algorithm1)consolidates the full TAKE pipeline\.

1

2

Input :corpus

𝒟=\{zn\}n=1N\\mathcal\{D\}=\\\{z\_\{n\}\\\}\_\{n=1\}^\{N\}; budget

KK; checkpoints

TT; kernel

k​\(⋅\)k\(\\cdot\);

ε\\varepsilon; pool size

MM
Output :distilled corpus

𝒟~\\tilde\{\\mathcal\{D\}\},

\|𝒟~\|=K\|\\tilde\{\\mathcal\{D\}\}\|=K
3

4

⊳\\trianglerightStage 1: Score

5Train logistic probe on

𝒟\\mathcal\{D\}, saving checkpoints

θ0,…,θT−1\\theta\_\{0\},\\ldots,\\theta\_\{T\-1\}
6for*t=0,…,T−1t=0,\\ldots,T\-1*do

7Compute

I\(n,t\)←‖∇Wℓ​\(zn;θt\)‖I\_\{\(n,t\)\}\\leftarrow\\\|\\nabla\_\{W\}\\ell\(z\_\{n\};\\,\\theta\_\{t\}\)\\\|for all

nn\(Goodfellow trick\)

8Normalize

I\(⋅,t\)←I\(⋅,t\)/∑n′I\(n′,t\)I\_\{\(\\cdot,t\)\}\\leftarrow I\_\{\(\\cdot,t\)\}/\\sum\_\{n^\{\\prime\}\}I\_\{\(n^\{\\prime\},t\)\}\(column\-wise\)

9Compute

κn←∑t=0T−1k​\(t\)/I\(n,t\)\\kappa\_\{n\}\\leftarrow\\textstyle\\sum\_\{t=0\}^\{T\-1\}k\(t\)/I\_\{\(n,t\)\}for all

nnvia Eq\.[9](https://arxiv.org/html/2607.11898#S4.E9)

10Normalize

wn←κn/∑n′κn′w\_\{n\}\\leftarrow\\kappa\_\{n\}/\\sum\_\{n^\{\\prime\}\}\\kappa\_\{n^\{\\prime\}\}for all

nn\(definesPwP^\{w\}\)

⊳\\trianglerightStage 2: Generate

11Fine\-tune small language model

LM\\mathrm\{LM\}on

𝒟\\mathcal\{D\}
12Sample synthetic pool

𝒟synth=\{z~c\}c=1M\\mathcal\{D\}\_\{\\text\{synth\}\}=\\\{\\tilde\{z\}\_\{c\}\\\}\_\{c=1\}^\{M\}from

LM\\mathrm\{LM\}
⊳\\trianglerightStage 3: Distill

13Embed:

ϕ​\(zn\)\\phi\(z\_\{n\}\)for

n∈\[N\]n\\in\[N\];

ϕ​\(z~c\)\\phi\(\\tilde\{z\}\_\{c\}\)for

c∈\[M\]c\\in\[M\]
14Compute cost matrix

Cn​c=‖ϕ​\(zn\)−ϕ​\(z~c\)‖2C\_\{nc\}=\\\|\\phi\(z\_\{n\}\)\-\\phi\(\\tilde\{z\}\_\{c\}\)\\\|^\{2\}
15Solve

γ∗←Sinkhorn​\(C,w,1/M,ε\)\\gamma^\{\*\}\\leftarrow\\mathrm\{Sinkhorn\}\(C,\\,w,\\,\\mathbf\{1\}/M,\\,\\varepsilon\)via Eq\.[11](https://arxiv.org/html/2607.11898#S4.E11)

16Distill

𝒟~←top​\-​K​\(\{z~c\},∑nγn​c∗\)\\tilde\{\\mathcal\{D\}\}\\leftarrow\\mathrm\{top\}\\text\{\-\}K\\\!\\left\(\\\{\\tilde\{z\}\_\{c\}\\\},\\;\\textstyle\\sum\_\{n\}\\gamma^\{\*\}\_\{nc\}\\right\)\(by received mass\)

return*𝒟~\\tilde\{\\mathcal\{D\}\}*

Algorithm 1Trajectory\-Aware Knowledge Estimation \(TAKE\)

## 5Experiments

### 5\.1Setup

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

- •Random: uniform random selection at the target budget\.
- •EDA\[[30](https://arxiv.org/html/2607.11898#bib.bib30)\]: word\-level augmentation \(10×\\times\) applied to theRandomset\.
- •DiLM\[[7](https://arxiv.org/html/2607.11898#bib.bib7)\]: LLM\-generated candidates selected by gradient matching across k\-centre selection\. It evaluated on both classification and NLI\.
- •DaLLME\[[8](https://arxiv.org/html/2607.11898#bib.bib8)\]: 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](https://arxiv.org/html/2607.11898#bib.bib31)\]features, TextCNN\[[32](https://arxiv.org/html/2607.11898#bib.bib32)\]and TextRNN\[[33](https://arxiv.org/html/2607.11898#bib.bib33)\]with GloVe\[[34](https://arxiv.org/html/2607.11898#bib.bib34)\]word embeddings, and BERT\[[35](https://arxiv.org/html/2607.11898#bib.bib35)\]\. 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±s\.e\.\{\}\_\{\\pm\{\\text\{s\.e\.\}\}\}across five runs\.

In terms of implementation, TAKE extracts knowledge scoresκn\\kappa\_\{n\}by runningT=20T=20training checkpoints with exponential kernelλ∗=T/Tm\\lambda^\{\*\}=T/T\_\{m\}, using all\-MiniLM\-L6\-v2\[[36](https://arxiv.org/html/2607.11898#bib.bib36)\]as the frozen encoderϕ\\phi\. The synthetic pool uses Gemma\-3\[[37](https://arxiv.org/html/2607.11898#bib.bib37)\]\(270 M\) fine\-tuned on𝒟\\mathcal\{D\}to generate\|𝒟synth\|=5,000\|\\mathcal\{D\}\_\{\\text\{synth\}\}\|=5\{,\}000samples\. OT selection embeds via the sameϕ\\phiand uses Sinkhorn withε=0\.05\\varepsilon=0\.05for 200 iterations\. All experiments fit on a single NVIDIA V100 \(40 GB\)\.

### 5\.2Classification

Table 2:Classification accuracy \(%\) on AG News, IMDb, and SST\-2\.MethodBudgetLRTextCNNTextRNNBERTAG NewsFull dataset100%90\.21±0\.1491\.27±0\.2090\.98±0\.3293\.85±0\.12Random0\.1%66\.39±2\.2775\.56±0\.1474\.91±1\.4578\.12±3\.25EDA0\.1%68\.03±2\.4382\.27±0\.1581\.46±1\.3484\.45±1\.44DaLLME0\.1%74\.60±0\.02†88\.30±0\.02†84\.50±0\.02†−\-TAKE0\.1%77\.07±0\.0988\.53±0\.2184\.25±1\.0389\.82±0\.18IMDbFull dataset100%88\.41±0\.1387\.40±0\.2083\.00±0\.3591\.22±0\.25Random0\.1%59\.73±5\.5157\.65±1\.1955\.09±4\.3160\.40±4\.10EDA0\.1%61\.52±4\.6262\.40±3\.4058\.92±1\.9363\.12±1\.95DaLLME0\.1%65\.00±0\.02†68\.30±0\.02†64\.50±0\.02†−\-TAKE0\.1%66\.83±0\.1368\.93±2\.1563\.26±0\.3871\.55±0\.22SST\-2Full dataset100%80\.23±0\.2788\.81±0\.1587\.40±0\.7692\.52±0\.28Random20/cls55\.28±2\.3263\.85±0\.8761\.52±1\.6669\.97±3\.21EDA20/cls57\.21±1\.0565\.54±0\.4563\.17±0\.8474\.21±1\.58DiLM20/cls−\-−\-−\-80\.30±2\.80†TAKE20/cls65\.05±0\.0570\.65±0\.3267\.18±0\.4981\.15±0\.24
- •\[−\-\] not reported; \[†\\dagger\] reported from respective paper; \[bold\] best per budget\.DaLLME reported with its best case \(OpenAI\-3\-large embeddings\)\.

Accuracy gains and their source\.TAKE \(0\.1%\) outperforms DaLLME on LR by\+2\.5\+2\.5on AG News and\+1\.8\+1\.8on 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 underPwP^\{w\}via Theorem[1](https://arxiv.org/html/2607.11898#Thmtheorem1)\(the gap to the originalP𝒟P\_\{\\mathcal\{D\}\}is bounded by Remark[1](https://arxiv.org/html/2607.11898#Thmremark1)\)\. Sinceκn\\kappa\_\{n\}scores are computed independently of the evaluation backbone, this coverage ofPwP^\{w\}transfers across model families\.

Stability across runs\.The wide standard errors for Random \(IMDb LR±5\.5\\pm 5\.5, BERT±4\.1\\pm 4\.1\) reveal that random selection at 0\.1% can land anywhere from a coherent corpus to a degenerate one\. TAKE’s errors are at least5×5\\timesnarrower \(±0\.1\\pm 0\.1–2\.22\.2\), a direct consequence of OT\-based selection: the transport plan distributes selected prototypes acrossPwP^\{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\.1581\.15vs\.80\.3080\.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\.3Natural Language Inference

Table 3:NLI accuracy \(%\) on MNLI\-m and QQP\.MethodBudgetSiamese LRBERTMNLI\-mFull dataset100%60\.02±0\.0786\.81±0\.30Random20/cls34\.15±2\.1040\.10±3\.20EDA20/cls37\.80±1\.4544\.52±2\.80DiLM20/cls−\-48\.70±2\.60†TAKE20/cls42\.66±0\.1551\.24±0\.45QNLIFull dataset100%73\.05±0\.1382\.50±0\.12Random20/cls52\.21±0\.2857\.19±2\.20EDA20/cls54\.33±1\.2159\.81±1\.42DiLM20/cls−\-−\-TAKE20/cls57\.96±0\.5864\.89±3\.58QQPFull dataset100%78\.58±0\.0589\.45±0\.15Random20/cls52\.40±1\.3259\.10±3\.80EDA20/cls55\.12±1\.8862\.35±2\.45DiLM20/cls−\-64\.40±2\.20†TAKE20/cls60\.04±0\.1267\.76±0\.35
- •\[−\-\] not reported; \[†\\dagger\] reported from respective paper;\[bold\] best per budget\.

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\+2\.5on MNLI\-m and\+3\.4\+3\.4on QQP on a BERT learner, and provides competitive results for QNLI\. The LR gap \(\+8\.5\+8\.5on MNLI\-m,\+5\.8\+5\.8on QNLI,\+7\.6\+7\.6on 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\.0460\.04vs\.59\.1059\.10\) and QNLI \(57\.9657\.96vs\.57\.1957\.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\.886\.8/82\.582\.5/89\.589\.5for 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 \(≈60\\approx\\\!60pairs for MNLI\-m,≈40\\approx\\\!40for 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 \(±0\.1\\pm 0\.1–0\.50\.5\) are55–7×7\\timesnarrower than DiLM’s \(±2\.2\\pm 2\.2–2\.62\.6\), mirroring the pattern on classification\. This is consistent with the OT transport plan enforcing coverage ofPwP^\{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\.4Component Contributions

Table 4:Distillation accuracy \(%\) under TAKE component ablation\.AG NewsQQPLRBERTSiamese LRBERTRandom66\.39±2\.2775\.92±0\.9852\.40±1\.3257\.94±1\.30▲kk\-Means nearest72\.89±1\.3777\.79±1\.0455\.28±1\.1261\.30±1\.08TAKE withoutκ\\kappa75\.82±0\.2380\.94±0\.9358\.80±0\.9263\.66±1\.25TAKE77\.07±0\.0986\.50±1\.0160\.04±0\.1266\.29±1\.91
- •\[▲\\blacktriangle\]notinclusive of TAKE; \[bold\] best results\.

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[4](https://arxiv.org/html/2607.11898#S5.T4)\. TAKE improves over Random by\+10\.7\+10\.7on AG News LR and\+10\.6\+10\.6on BERT;kk\-means, by comparison, accounts for only part of this gain \(\+6\.5\+6\.5LR,\+1\.9\+1\.9BERT\), while our discrete OT selection step alone \(TAKE withoutκ\\kappa\) already achieves\+9\.4\+9\.4LR and\+5\.0\+5\.0BERT from the same baseline\. The same pattern holds on QQP:kk\-means gains\+2\.9\+2\.9/\+3\.4\+3\.4\(Siamese LR / BERT\) whereas TAKE withoutκ\\kappagains\+6\.4\+6\.4/\+5\.7\+5\.7\. Addingκ\\kappaknowledge scores delivers the remaining improvement, modestly for Siamese LR \(\+1\.2\+1\.2\) but substantially for BERT \(\+5\.6\+5\.6on AG News,\+2\.6\+2\.6on QQP\), suggesting thatκ\\kappaprimarily benefits models with enough capacity to exploit the harder, informative examples that high\-κ\\kappasamples represent\.

### 5\.5Kernels for Temporal Weighting

Table 5:Distillation accuracy \(%\) with alternative temporal kernels\.AG NewsQQPKernelLRBERTSiamese LRBERTTAKE withoutκ\\kappa75\.82±0\.2380\.94±0\.9358\.80±0\.9263\.66±1\.25Last checkpoint75\.24±0\.5778\.28±0\.6658\.32±1\.5662\.19±0\.57Constant76\.13±1\.1483\.47±0\.7159\.44±0\.8664\.78±1\.51Linear76\.20±0\.6885\.20±1\.0059\.61±1\.1165\.60±1\.35Cosine76\.74±0\.5385\.46±0\.6459\.14±0\.8566\.16±0\.41Exponential \(TAKE\)77\.07±0\.0986\.50±1\.0160\.04±0\.1266\.29±1\.91
- •\[bold\] best results\.

Table[5](https://arxiv.org/html/2607.11898#S5.T5)ablates the choice of temporal kernel defined in §[4\.2\.2](https://arxiv.org/html/2607.11898#S4.SS2.SSS2)\. The last\-checkpoint baseline is the critical exception: it fallsbelowthe no\-κ\\kappabaseline across all backbones, most severely on BERT \(−2\.7\-2\.7on AG News,−1\.5\-1\.5on QQP\)\. This is precisely the failure mode identified in Proposition[2](https://arxiv.org/html/2607.11898#Thmproposition2): whenκ\\kappais extracted at convergence only, the learner has fully memorized difficult samples, so thehard\-sample biasis 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\-κ\\kappabaseline \(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\.

## 6Limitations & Ethical Considerations

We identify three limitations of the current method\. \(1\)Gradient norm proxy\.I\(n,t\)I\_\{\(n,t\)\}uses plain gradient norms rather than Fisher\-preconditioned scores; the Goodfellow trick\[[28](https://arxiv.org/html/2607.11898#bib.bib28)\]keeps per\-sample cost atO​\(C\+dϕ\)O\(C\+d\_\{\\phi\}\)per checkpoint\. K\-FAC\[[38](https://arxiv.org/html/2607.11898#bib.bib38)\]or EKFAC\[[39](https://arxiv.org/html/2607.11898#bib.bib39)\]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\. TheLL\-Lipschitz assumption onℓ\\ell\(Theorem 1\) holds approximately whenϕ\\phiis 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𝒟train\\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𝒟train\\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𝒟synth\\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\.

## 7Conclusion

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\.

## References

- \[1\]Tongzhou Wang, Jun\-Yan Zhu, Antonio Torralba, and Alexei A Efros\.Dataset distillation\.arXiv preprint arXiv:1811\.10959, 2018\.
- \[2\]Bo Zhao, Konda Reddy Mopuri, and Hakan Bilen\.Dataset condensation with gradient matching\.arXiv preprint arXiv:2006\.05929, 2020\.
- \[3\]George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei A Efros, and Jun\-Yan Zhu\.Dataset distillation by matching training trajectories\.InCVPR, 2022\.
- \[4\]Ilia Sucholutsky and Matthias Schonlau\.Soft\-label dataset distillation and text dataset distillation\.InIJCNN, 2021\.
- \[5\]Yongqi Li and Wenjie Li\.Data distillation for text classification\.arXiv preprint arXiv:2104\.08448, 2021\.
- \[6\]Aru Maekawa, Naoki Kobayashi, Kotaro Funakoshi, and Manabu Okumura\.Dataset distillation with attention labels for fine\-tuning bert\.InACL, 2023\.
- \[7\]Aru Maekawa, Satoshi Kosugi, Kotaro Funakoshi, and Manabu Okumura\.Dilm: Distilling dataset into language model for text\-level dataset distillation\.InNAACL Findings, 2024\.
- \[8\]Yefan Tao, Luyang Kong, Andrey Kan, and Laurent Callot\.Textual dataset distillation via language model embedding\.InEMNLP Findings, 2024\.
- \[9\]Pang Wei Koh and Percy Liang\.Understanding black\-box predictions via influence functions\.InICML, 2017\.
- \[10\]Kai Wang, Bo Zhao, Xiangyu Peng, Zheng Zhu, Shuo Yang, Shuo Wang, Guan Huang, Hakan Bilen, Xinchao Wang, and Yang You\.Cafe: Learning to condense dataset by aligning features\.InCVPR, 2022\.
- \[11\]Jang\-Hyun Kim, Jinuk Kim, Seong Joon Oh, Sangdoo Yun, Hwanjun Song, Joonhyun Jeong, Jung\-Woo Ha, and Hyun Oh Song\.Dataset condensation via efficient synthetic\-data parameterization\.InICML, 2022\.
- \[12\]R Dennis Cook and Sanford Weisberg\.Residuals and influence in regression\.New York: Chapman and Hall, 1982\.
- \[13\]Garima Pruthi, Frederick Liu, Satyen Kale, and Mukund Sundararajan\.Estimating training data influence by tracing gradient descent\.NeurIPS, 2020\.
- \[14\]Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry\.Trak: Attributing model behavior at scale\.InICML, 2023\.
- \[15\]Yongchan Kwon, Eric Wu, Kevin Wu, and James Zou\.Datainf: Efficiently estimating data influence in lora\-tuned llms and diffusion models\.InICLR, 2023\.
- \[16\]Devansh Arpit, Stanisław Jastrzębski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S Kanwal, Tegan Maharaj, Asja Fischer, Aaron Courville, Yoshua Bengio, et al\.A closer look at memorization in deep networks\.InICML, 2017\.
- \[17\]Swabha Swayamdipta, Roy Schwartz, Nicholas Lourie, Yizhong Wang, Hannaneh Hajishirzi, Noah A Smith, and Yejin Choi\.Dataset cartography: Mapping and diagnosing datasets with training dynamics\.InEMNLP, 2020\.
- \[18\]Bo Han, Quanming Yao, Xingrui Yu, Gang Niu, Miao Xu, Weihua Hu, Ivor Tsang, and Masashi Sugiyama\.Co\-teaching: Robust training of deep neural networks with extremely noisy labels\.NeurIPS, 2018\.
- \[19\]Bo Zhao and Hakan Bilen\.Dataset condensation with distribution matching\.InWACV, 2023\.
- \[20\]Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola\.A kernel two\-sample test\.The journal of machine learning research, 13\(1\):723–773, 2012\.
- \[21\]Liam Paninski\.Estimation of entropy and mutual information\.Neural computation, 15\(6\):1191–1253, 2003\.
- \[22\]Marco Cuturi\.Sinkhorn distances: Lightspeed computation of optimal transport\.NeurIPS, 2013\.
- \[23\]Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger\.From word embeddings to document distances\.InICML, 2015\.
- \[24\]Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf\.Wasserstein auto\-encoders\.InICLR, 2018\.
- \[25\]Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy\.Optimal transport for domain adaptation\.PAMI, 2016\.
- \[26\]Martin Arjovsky, Soumith Chintala, and Léon Bottou\.Wasserstein generative adversarial networks\.InICML, 2017\.
- \[27\]Han Guo, Nazneen Rajani, Peter Hase, Mohit Bansal, and Caiming Xiong\.Fastif: Scalable influence functions for efficient model interpretation and debugging\.InProceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10333–10350, 2021\.
- \[28\]Ian Goodfellow\.Efficient per\-example gradient computations\.arXiv preprint arXiv:1510\.01799, 2015\.
- \[29\]Olivier Bousquet and André Elisseeff\.Stability and generalization\.JMLR, 2002\.
- \[30\]Jason Wei and Kai Zou\.Eda: Easy data augmentation techniques for boosting performance on text classification tasks\.InProceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference on natural language processing \(EMNLP\-IJCNLP\), pages 6382–6388, 2019\.
- \[31\]Gerard Salton and Christopher Buckley\.Term\-weighting approaches in automatic text retrieval\.Information processing & management, 24\(5\):513–523, 1988\.
- \[32\]Yoon Kim\.Convolutional neural networks for sentence classification\.InEMNLP, 2014\.
- \[33\]Pengfei Liu, Xipeng Qiu, and Xuanjing Huang\.Recurrent neural network for text classification with multi\-task learning\.arXiv preprint arXiv:1605\.05101, 2016\.
- \[34\]Jeffrey Pennington, Richard Socher, and Christopher D Manning\.Glove: Global vectors for word representation\.InProceedings of the 2014 conference on empirical methods in natural language processing \(EMNLP\), pages 1532–1543, 2014\.
- \[35\]Jacob Devlin, Ming\-Wei Chang, Kenton Lee, and Kristina Toutanova\.Bert: Pre\-training of deep bidirectional transformers for language understanding\.InProceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 \(long and short papers\), pages 4171–4186, 2019\.
- \[36\]Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou\.Minilm: Deep self\-attention distillation for task\-agnostic compression of pre\-trained transformers\.Advances in neural information processing systems, 33:5776–5788, 2020\.
- \[37\]Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, Louis Rouillard, et al\.Gemma 3 technical report\.arXiv preprint arXiv:2503\.19786, 4, 2025\.
- \[38\]James Martens and Roger Grosse\.Optimizing neural networks with kronecker\-factored approximate curvature\.InICML, 2015\.
- \[39\]Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent\.Fast approximate natural gradient descent in a kronecker factored eigenbasis\.NeurIPS, 2018\.
- \[40\]Cédric Villani et al\.Optimal transport: old and new, volume 338\.Springer, 2009\.
- \[41\]Nils Reimers and Iryna Gurevych\.Sentence\-bert: Sentence embeddings using siamese bert\-networks\.InProceedings of the 2019 conference on empirical methods in natural language processing and the 9th international joint conference on natural language processing \(EMNLP\-IJCNLP\), pages 3982–3992, 2019\.

\[Supplementary Material\]

TAKE: Trajectory\-Aware Knowledge Estimation for Text Dataset Distillation

## Appendix AProofs

### A\.1Proof of Proposition 1 \(Hard\-Sample Bias\)

Setup\.Letℐself​\(z\)=∇⊤Hθ^−1​∇\\mathcal\{I\}\_\{\\text\{self\}\}\(z\)=\\nabla^\{\\top\}H\_\{\\hat\{\\theta\}\}^\{\-1\}\\nabla, where∇≡∇θℓ​\(z;θ^\)\\nabla\\equiv\\nabla\_\{\\theta\}\\ell\(z;\\hat\{\\theta\}\)andHθ^H\_\{\\hat\{\\theta\}\}is the empirical Hessian, assumed symmetric positive definite \(SPD\) withλmin​\(Hθ^\)\>0\\lambda\_\{\\min\}\(H\_\{\\hat\{\\theta\}\}\)\>0\(this holds whenever the model is not at a degenerate critical point\)\. Partition𝒟\\mathcal\{D\}into𝒟clean\\mathcal\{D\}\_\{\\text\{clean\}\}and𝒟noisy\\mathcal\{D\}\_\{\\text\{noisy\}\}such that‖∇‖≤ϵ\\\|\\nabla\\\|\\leq\\epsilonfor allz∈𝒟cleanz\\in\\mathcal\{D\}\_\{\\text\{clean\}\}and‖∇‖≥δ\\\|\\nabla\\\|\\geq\\deltafor allz∈𝒟noisyz\\in\\mathcal\{D\}\_\{\\text\{noisy\}\}, withδ≫ϵ\\delta\\gg\\epsilon\.

Step 1 — Boundℐself\\mathcal\{I\}\_\{\\text\{self\}\}for each partition\.SinceHθ^H\_\{\\hat\{\\theta\}\}is SPD with eigenvalues in\[λmin​\(Hθ^\),λmax​\(Hθ^\)\]\[\\lambda\_\{\\min\}\(H\_\{\\hat\{\\theta\}\}\),\\,\\lambda\_\{\\max\}\(H\_\{\\hat\{\\theta\}\}\)\], the Rayleigh quotient gives:

‖∇‖2λmax​\(Hθ^\)≤ℐself​\(z\)≤‖∇‖2λmin​\(Hθ^\)\.\\frac\{\\\|\\nabla\\\|^\{2\}\}\{\\lambda\_\{\\max\}\(H\_\{\\hat\{\\theta\}\}\)\}\\;\\leq\\;\\mathcal\{I\}\_\{\\text\{self\}\}\(z\)\\;\\leq\\;\\frac\{\\\|\\nabla\\\|^\{2\}\}\{\\lambda\_\{\\min\}\(H\_\{\\hat\{\\theta\}\}\)\}\.\(13\)Substituting the gradient\-norm assumptions \(‖∇‖2≤ϵ2\\\|\\nabla\\\|^\{2\}\\leq\\epsilon^\{2\}for clean samples;‖∇‖2≥δ2\\\|\\nabla\\\|^\{2\}\\geq\\delta^\{2\}for noisy samples\) and taking expectations \(the bounds are uniform over each partition, so expectation preserves them\):

𝔼z∼𝒟clean​\[ℐself​\(z\)\]\\displaystyle\\mathbb\{E\}\_\{z\\sim\\mathcal\{D\}\_\{\\text\{clean\}\}\}\[\\mathcal\{I\}\_\{\\text\{self\}\}\(z\)\]≤ϵ2λmin​\(Hθ^\),\\displaystyle\\leq\\frac\{\\epsilon^\{2\}\}\{\\lambda\_\{\\min\}\(H\_\{\\hat\{\\theta\}\}\)\},\(14\)𝔼z∼𝒟noisy​\[ℐself​\(z\)\]\\displaystyle\\mathbb\{E\}\_\{z\\sim\\mathcal\{D\}\_\{\\text\{noisy\}\}\}\[\\mathcal\{I\}\_\{\\text\{self\}\}\(z\)\]≥δ2λmax​\(Hθ^\)\.\\displaystyle\\geq\\frac\{\\delta^\{2\}\}\{\\lambda\_\{\\max\}\(H\_\{\\hat\{\\theta\}\}\)\}\.\(15\)
Step 2 — Lower bound onΔHSB\\Delta\_\{\\text\{HSB\}\}\.

ΔHSB:=𝔼z∼𝒟noisy​\[ℐself​\(z;θ^\)\]−𝔼z∼𝒟clean​\[ℐself​\(z;θ^\)\]≥δ2λmax​\(Hθ^\)−ϵ2λmin​\(Hθ^\)\.\\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\}\}\)\}\.\(16\)Substitutingλmin​\(Hθ^\)=λmax​\(Hθ^\)/cond​\(Hθ^\)\\lambda\_\{\\min\}\(H\_\{\\hat\{\\theta\}\}\)=\\lambda\_\{\\max\}\(H\_\{\\hat\{\\theta\}\}\)/\\mathrm\{cond\}\(H\_\{\\hat\{\\theta\}\}\):

ΔHSB≥δ2−cond​\(Hθ^\)​ϵ2λmax​\(Hθ^\)\>0wheneverδϵ\>cond​\(Hθ^\)\.\\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\}\}\)\}\.\(17\)Under the interpolation regime \(an overparameterised model trained to near\-zero training loss\),ϵ→0\\epsilon\\to 0: the bound simplifies toδ2/λmax​\(Hθ^\)\>0\\delta^\{2\}/\\lambda\_\{\\max\}\(H\_\{\\hat\{\\theta\}\}\)\>0and the positivity condition is trivially satisfied \(δ/ϵ→∞\\delta/\\epsilon\\to\\infty\)\.□\\square

### A\.2Proof of Proposition 2 \(Trajectory Integration Reduces HSB\)

Proposition 2\.Under smooth loss and bounded gradient noise, the trajectory\-integrated knowledge scoreκn\\kappa\_\{n\}produces a strictly smaller noisy\-over\-clean gap than the single last\-checkpoint baselineκn†=1/I\(n,T−1\)\\kappa^\{\\dagger\}\_\{n\}=1/I\_\{\(n,T\-1\)\}:

𝔼z∼𝒟noisy​\[κn\]−𝔼z∼𝒟clean​\[κn\]<𝔼z∼𝒟noisy​\[κn†\]−𝔼z∼𝒟clean​\[κn†\]\.\\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\}\]\.\(18\)
Memorization onsetTmT\_\{m\}\.In the main paper \(§4\.2\.2\) we defineTmT\_\{m\}as the first checkpoint at which noisy\-sample gradient norms begin to decrease\. The formal continuous analogue is:

Tm:=inf\{t:𝔼z∼𝒟noisy​\[‖∇θℓ​\(z;θt\)‖\]<𝔼z∼𝒟noisy​\[‖∇θℓ​\(z;θ0\)‖\]\}\.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\\\}\.\(19\)Assumptions\.

- \(A1\)Monotone gradient decay for clean samples\.Forz∈𝒟cleanz\\in\\mathcal\{D\}\_\{\\text\{clean\}\},I\(n,t\)I\_\{\(n,t\)\}is non\-increasing intt, converging toward zero \(holds underLL\-smooth loss and a sufficiently small learning rate\)\. The reciprocal1/I\(n,t\)1/I\_\{\(n,t\)\}is therefore non\-decreasing\.
- \(A2\)Bounded gradient for noisy samples within the proof horizon\.Forz∈𝒟noisyz\\in\\mathcal\{D\}\_\{\\text\{noisy\}\},I\(n,t\)≥δ\>0I\_\{\(n,t\)\}\\geq\\delta\>0for allt≤Tmt\\leq T\_\{m\}, whereTmT\_\{m\}is the memorization onset defined in Section 4\.2\.2 of the main paper\. BeyondTmT\_\{m\}, noisy gradients decay; the decreasing kernel discounts these late checkpoints\. Withλ∗=T/Tm\\lambda^\{\*\}=T/T\_\{m\}, the kernel weight atTm\+1T\_\{m\}\+1relative tot∗t^\{\*\}ise−λ∗​\(Tm\+1−t∗\)/Te^\{\-\\lambda^\{\*\}\(T\_\{m\}\+1\-t^\{\*\}\)/T\}, which decreases asTm→TT\_\{m\}\\to T, making the post\-memorization contributionS2S\_\{2\}negligible relative toS1S\_\{1\}\.
- \(A3\)Post\-memorization baseline inflation\.Att=T−1\>Tmt=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:ΔT−1:=𝔼noisy​\[1/I\(n,T−1\)\]−𝔼clean​\[1/I\(n,T−1\)\]\>0\\Delta\_\{T\-1\}:=\\mathbb\{E\}\_\{\\text\{noisy\}\}\[1/I\_\{\(n,T\-1\)\}\]\-\\mathbb\{E\}\_\{\\text\{clean\}\}\[1/I\_\{\(n,T\-1\)\}\]\>0\.
- \(A4\)Positive temporal kernel\.k​\(t\)\>0k\(t\)\>0for alltt; decreasing sok​\(t∗\)\>k​\(T−1\)k\(t^\{\*\}\)\>k\(T\-1\)for anyt∗<T−1t^\{\*\}<T\-1\.

Proof\.

DefineΔt:=𝔼noisy​\[1/I\(n,t\)\]−𝔼clean​\[1/I\(n,t\)\]\\Delta\_\{t\}:=\\mathbb\{E\}\_\{\\text\{noisy\}\}\[1/I\_\{\(n,t\)\}\]\-\\mathbb\{E\}\_\{\\text\{clean\}\}\[1/I\_\{\(n,t\)\}\], so

𝔼noisy​\[κn\]−𝔼clean​\[κn\]=∑t=0T−1k​\(t\)​Δt\.\\mathbb\{E\}\_\{\\text\{noisy\}\}\[\\kappa\_\{n\}\]\-\\mathbb\{E\}\_\{\\text\{clean\}\}\[\\kappa\_\{n\}\]=\\sum\_\{t=0\}^\{T\-1\}k\(t\)\\,\\Delta\_\{t\}\.\(20\)
Step 1 — Sign structure\.For1≤t≤Tm1\\leq t\\leq T\_\{m\}: \(A1\) and \(A2\) together giveΔt≤0\\Delta\_\{t\}\\leq 0, strictly negative at somet∗≤Tmt^\{\*\}\\leq T\_\{m\}\. Fort\>Tmt\>T\_\{m\}:Δt\\Delta\_\{t\}may be positive \(memorization onset decreases noisy norms\)\.

Step 2 — Weighted sum\.Split atTmT\_\{m\}:

S1=∑t=1Tmk​\(t\)​Δt<0andS2=∑t=Tm\+1T−1k​\(t\)​Δt≥0\.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\.\(21\)Since the kernel is decreasing,S2≤k​\(Tm\+1\)⋅∑t\>TmΔtS\_\{2\}\\leq k\(T\_\{m\}\+1\)\\cdot\\sum\_\{t\>T\_\{m\}\}\\Delta\_\{t\}, which is dominated by\|S1\|\|S\_\{1\}\|for sufficiently largeλ\\lambda\. Hence∑tk​\(t\)​Δt<0\\sum\_\{t\}k\(t\)\\,\\Delta\_\{t\}<0\.

Step 3 — Baseline gap\.By \(A3\),𝔼noisy​\[κn†\]−𝔼clean​\[κn†\]=ΔT−1\>0\\mathbb\{E\}\_\{\\text\{noisy\}\}\[\\kappa^\{\\dagger\}\_\{n\}\]\-\\mathbb\{E\}\_\{\\text\{clean\}\}\[\\kappa^\{\\dagger\}\_\{n\}\]=\\Delta\_\{T\-1\}\>0\.

Combining\.

𝔼noisy\[κn\]−𝔼clean\[κn\]<0<ΔT−1=𝔼noisy\[κn†\]−𝔼clean\[κn†\]\.□\\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\(22\)

### A\.3Distribution Matching Gap Bound \(Theorem 1\)

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

Goal\.Letf𝒮∗f^\{\*\}\_\{\\mathcal\{S\}\}be the model trained on𝒮\\mathcal\{S\}andℒ​\(𝒮\)=𝔼z∼Pw​\[ℓ​\(z;f𝒮∗\)\]\\mathcal\{L\}\(\\mathcal\{S\}\)=\\mathbb\{E\}\_\{z\\sim P^\{w\}\}\[\\ell\(z;f^\{\*\}\_\{\\mathcal\{S\}\}\)\]be the downstream loss underPwP^\{w\}\. Show:

\|ℒ​\(𝒟~\)−ℒ​\(𝒟\)\|≤2​L⋅W1​\(Pw,P𝒟~\)\+βK\.\|\\mathcal\{L\}\(\\tilde\{\\mathcal\{D\}\}\)\-\\mathcal\{L\}\(\\mathcal\{D\}\)\|\\;\\leq\\;2L\\cdot W\_\{1\}\(P^\{w\},P\_\{\\tilde\{\\mathcal\{D\}\}\}\)\+\\beta\_\{K\}\.\(23\)
Assumptions\.

- \(B1\)ℓ​\(⋅;f\)\\ell\(\\cdot;\\,f\)isLL\-Lipschitz in the embedding metric‖ϕ​\(⋅\)−ϕ​\(⋅\)‖\\\|\\phi\(\\cdot\)\-\\phi\(\\cdot\)\\\|for allff\.
- \(B2\)The model class has uniform stability constantβK=O​\(1/K\)\\beta\_\{K\}=O\(1/K\)\[[29](https://arxiv.org/html/2607.11898#bib.bib29)\]:supz\|ℓ​\(z;f𝒟~∗\)−ℓ​\(z;f𝒟∗\)\|≤βK\\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; asK→∞K\\to\\infty,βK→0\\beta\_\{K\}\\to 0\.
- \(B3\)The ground metric forW1W\_\{1\}isc​\(z,z′\)=‖ϕ​\(z\)−ϕ​\(z′\)‖c\(z,z^\{\\prime\}\)=\\\|\\phi\(z\)\-\\phi\(z^\{\\prime\}\)\\\|\.

Step 1 — Kantorovich–Rubinstein bound for a fixed model\.By Kantorovich–Rubinstein duality\[[40](https://arxiv.org/html/2607.11898#bib.bib40)\]under \(B3\), for any fixedff:

\|𝔼z∼Pw​\[ℓ​\(z;f\)\]−𝔼z∼P𝒟~​\[ℓ​\(z;f\)\]\|≤L⋅W1​\(Pw,P𝒟~\)\.\\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\}\}\}\)\.\(24\)
Step 2 — Triangle inequality decomposition\.ℒ​\(𝒟~\)\\mathcal\{L\}\(\\tilde\{\\mathcal\{D\}\}\)usesf𝒟~∗f^\{\*\}\_\{\\tilde\{\\mathcal\{D\}\}\}whileℒ​\(𝒟\)\\mathcal\{L\}\(\\mathcal\{D\}\)usesf𝒟∗f^\{\*\}\_\{\\mathcal\{D\}\}; both are evaluated underPwP^\{w\}but with different models, so Step 1 does not apply directly\. Insert two intermediates and apply the triangle inequality twice:

\|ℒ​\(𝒟~\)−ℒ​\(𝒟\)\|\\displaystyle\|\\mathcal\{L\}\(\\tilde\{\\mathcal\{D\}\}\)\-\\mathcal\{L\}\(\\mathcal\{D\}\)\|≤\|𝔼Pw​\[ℓ​\(z;f𝒟~∗\)\]−𝔼P𝒟~​\[ℓ​\(z;f𝒟~∗\)\]\|⏟\(I\) distribution shift, fixed​f𝒟~∗\\displaystyle\\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\}\}\}\}\+\|𝔼P𝒟~​\[ℓ​\(z;f𝒟~∗\)\]−𝔼P𝒟~​\[ℓ​\(z;f𝒟∗\)\]\|⏟\(II\) model difference, fixed​P𝒟~\\displaystyle\\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\}\}\}\}\+\|𝔼P𝒟~​\[ℓ​\(z;f𝒟∗\)\]−𝔼Pw​\[ℓ​\(z;f𝒟∗\)\]\|⏟\(III\) distribution shift, fixed​f𝒟∗\.\\displaystyle\\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\}\}\}\.\(25\)Terms \(I\) and \(III\) each satisfy Step 1 withf=f𝒟~∗f=f^\{\*\}\_\{\\tilde\{\\mathcal\{D\}\}\}andf=f𝒟∗f=f^\{\*\}\_\{\\mathcal\{D\}\}respectively, so each is bounded byL⋅W1​\(Pw,P𝒟~\)L\\cdot W\_\{1\}\(P^\{w\},P\_\{\\tilde\{\\mathcal\{D\}\}\}\)\. Term \(II\) is bounded pointwise by \(B2\):\(II\)≤βK=O​\(1/K\)\\text\{\(II\)\}\\leq\\beta\_\{K\}=O\(1/K\)\.

Combining\.

\|ℒ​\(𝒟~\)−ℒ​\(𝒟\)\|≤2​L⋅W1​\(Pw,P𝒟~\)\+βK\.\|\\mathcal\{L\}\(\\tilde\{\\mathcal\{D\}\}\)\-\\mathcal\{L\}\(\\mathcal\{D\}\)\|\\;\\leq\\;2L\\cdot W\_\{1\}\(P^\{w\},P\_\{\\tilde\{\\mathcal\{D\}\}\}\)\+\\beta\_\{K\}\.\(26\)MinimizingW1​\(Pw,P𝒟~\)W\_\{1\}\(P^\{w\},P\_\{\\tilde\{\\mathcal\{D\}\}\}\)via the OT objective directly minimizes the dominant term;βK→0\\beta\_\{K\}\\to 0asK→∞K\\to\\infty, a guarantee that k\-means and greedy coreset methods lack\.□\\square

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

## Appendix BEntropic 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μ=Pw=∑nwn​δzn\\mu=P^\{w\}=\\sum\_\{n\}w\_\{n\}\\delta\_\{z\_\{n\}\}\(TAKE\-weighted source\) andν=1\|𝒟synth\|​∑cδz~c\\nu=\\frac\{1\}\{\|\\mathcal\{D\}\_\{\\text\{synth\}\}\|\}\\sum\_\{c\}\\delta\_\{\\tilde\{z\}\_\{c\}\}\(uniform target over the synthetic pool\)\.

Transportation cost \(Kantorovich form\)\.

𝒞​\(γ\)=∫c​\(x,x^\)​𝑑γ​\(x,x^\),c​\(x,x^\)=‖ϕ​\(x\)−ϕ​\(x^\)‖2,γ∈Π​\(μ,ν\)\.\\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\)\.\(27\)
Entropic regularization\.The bare Kantorovich problem isO​\(N3\)O\(N^\{3\}\)complexity and non\-differentiable\. Adding an entropic penalty yields a strictly convex, smooth objective:

minγ∈Π​\(μ,ν\)⁡⟨C,γ⟩\+ε⋅KL​\(γ,μ⊗ν\),KL​\(γ,μ⊗ν\)=∫log⁡d​γd​\(μ⊗ν\)​d​γ,\\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,\(28\)whereμ⊗ν\\mu\\otimes\\nuis the product \(independence\) reference coupling\. Smallerε\\varepsilonrecovers the exact OT solution; largerε\\varepsilonsmooths the plan toward the independent coupling, encouraging diversity in the selected prototypes\.

Optional prior coupling\.If a task\-specific priorα​\(x,x^\)\\alpha\(x,\\hat\{x\}\)is available \(e\.g\., a topic\-conditional measure\), replaceμ⊗ν\\mu\\otimes\\nuwithα\\alphato regularize toward domain knowledge\.

Sinkhorn–Knopp solver\.The solution is computed by alternating row/column normalization on the kernel matrix𝐊n​c=e−Cn​c/ε\\mathbf\{K\}\_\{nc\}=e^\{\-C\_\{nc\}/\\varepsilon\}, yieldingO​\(N⋅\|𝒟synth\|\)O\(N\\cdot\|\\mathcal\{D\}\_\{\\text\{synth\}\}\|\)per iteration\.

Prototype extraction\.The distilled set𝒟~\\tilde\{\\mathcal\{D\}\}is formed by selecting the top\-KKcandidates ranked by received mass∑nγn​c∗\\sum\_\{n\}\\gamma^\{\*\}\_\{nc\}\.

## Appendix CImplementation Details

Datasets\.Table[6](https://arxiv.org/html/2607.11898#A3.T6)reports the task type and label set for each benchmark\.

Table 6:Dataset overview\.DatasetTaskNNLabelsAG NewsClassification120,000World, Sports, Business, Sci/TechIMDbClassification25,000Pos\., Neg\.SST\-2Classification67,349Pos\., Neg\.MNLI\-mNLI392,702Entailment, Neutral, ContradictionQNLINLI104,743Entailment, Not EntailmentQQPNLI363,846Duplicate, Non\-duplicateBackbone 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\.

Table 7:Backbone model details\.ModelArchitectureTasksLogisticLinear head on embeddingsClassificationTextCNNConvolutional encoderClassificationTextRNNRecurrent encoder \(GRU\)ClassificationSiamese LogisticSiamese features & logistic headNLIBERTTransformer encoder \(110M\)Classification \+ NLIHyperparameters\.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\[[41](https://arxiv.org/html/2607.11898#bib.bib41),[36](https://arxiv.org/html/2607.11898#bib.bib36)\]as the frozen sentence encoderϕ\\phi\. We useT=20T=20checkpoints, balancing trajectory coverage and compute\. The kernel decay is set toλ∗=T/Tm\\lambda^\{\*\}=T/T\_\{m\}, concentrating mass on the pre\-memorization phase as derived in Section 4\.2\.2\. For the Sinkhorn solver we useε=0\.05\\varepsilon=0\.05with 200 iterations, which produces soft plans with sufficient diversity and converges for all tested pool sizes\. The synthetic pool size is fixed at\|𝒟synth\|=5,000\|\\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\.