DiLM: Distilling Dataset into Language Model for Text-level Dataset Distillation

What came before

All prior text dataset distillation (TDD/sldd-2021's meta-learning approach, and attention-label distillation) optimizes synthetic examples as sequences of continuous word embeddings rather than discrete text, because gradient-based distillation objectives (originally designed for pixel-continuous images, e.g. DC's gradient matching) need a differentiable input.

The problem

Embedding-level distilled datasets are tied to the exact embedding table of the model used for distillation and cannot train a different model; nearest-neighbor decoding of the embeddings back to words (tried by sldd-2021 and later work) produces unrelated, unreadable words, so embedding-level synthetic sets are neither model-agnostic nor interpretable.

The idea

Instead of optimizing the discrete text or its embeddings directly, train the continuous parameters of a small generator language model (GPT-2) to emit informative synthetic text, backpropagating the gradient-matching loss into the generator's weights via a differentiable path (loss weighting by each sample's generation probability) that bypasses the non-differentiable sampling step.

How it works

Three stages: (1) pretrain the generator LM with a standard language-modeling loss on the real training data, using class-specific `<bos_i>` tokens so it learns to generate class-conditioned text; (2) fine-tune the generator with a nested bi-level loop — an outer loop that reinitializes a learner model (BERT-base) from $p(\theta_0)$, an inner loop that updates the learner on real data for $K$ steps, and a per-inner-step DC-style gradient-matching loss $\mathcal{L}_{GM}=D(\nabla_\theta \mathcal{L}_{real},\nabla_\theta\mathcal{L}_{syn})$ (cosine distance) computed per class, where $\mathcal{L}_{syn}=\sum_i a_i\, l_\theta(\tilde{x}_i)$ weights each sampled synthetic example's learner loss by its normalized generation probability $a_i$ so gradients can flow into the generator $\phi$; two add-on techniques improve this — a "representative teacher" (K-center clustering of real samples used as the real-data mini-batch instead of random sampling) and "diverse mini-batch sampling" (K-means clustering of an oversampled batch of generated candidates so each training step sees diverse synthetic samples); (3) generate the final distilled dataset by sampling 100x the target size via top-$p$ sampling and selecting K-center representative samples. Only the randomly initialized last layer of the learner is used in gradient matching to reduce cost. The stored artifact is the trained generator LM (or, for evaluation comparability, a fixed sampled text set).

Evidence

Table 1 (SST-2/QQP/MNLI-m, BERT-base learner, hard labels, DPC=5/10/20): DiLM beats K-centers coreset selection at every setting (e.g. SST-2 DPC=20: 80.3% vs. 79.8%; MNLI-m DPC=20: 48.7% vs. 45.3%; full dataset 86.7%), and beats a "Vanilla LM" (generator without gradient-matching fine-tuning) by a wide margin (e.g. SST-2 DPC=5: 72.5% vs. 65.2%), isolating that gradient matching, not just LM-quality text generation, drives the gain. TDD (sldd-2021), evaluated as raw embeddings, beats DiLM and even the full dataset (SST-2 89.6%) but collapses to near-random (50.2%) once its embeddings are decoded to nearest- neighbor text — the paper's key demonstration that embedding-level distillation is not usable as text. Table 2 (cross-model generalization, DPC=20): DiLM-distilled text (generated using BERT-base) trains RoBERTa-base, BERT-large, and XLNet-base (a different architecture family) better than K-centers in 11/12 cells (e.g. SST-2 XLNet-base: 77.9% vs. 71.8%). Table 3 (5-shot in-context learning of SST-2): DiLM-selected examples beat Random and K-centers as ICL prompts across GPT-2-XL (1.5B), OPT (2.7B), and Llama 2 (7B) (e.g. Llama 2: 95.1% vs. 94.6% K-centers). Table 4 ablation (DPC=5): removing the representative-teacher, diverse-mini-batch-sampling, or final K-center selection component each individually drops accuracy (e.g. SST-2 drops from 72.5% to 65.2% without selection), confirming all three additions contribute.

Limitations

A performance gap to the full dataset remains (e.g. SST-2 DPC=20 80.3% vs. 86.7% full); only classification tasks are tested, not generation or full LM pretraining; the small GPT-2 (128M) generator shows visible repetition artifacts on the harder MNLI-m task; training the generator adds upfront cost beyond directly optimizing embeddings; the paper explicitly notes DiLM cannot be used for privacy preservation, since the generator LM can memorize and regenerate real training samples (unlike grad-mm-2025's later explicit DP mechanism for the same discrete-text problem).