TGDD: Trajectory Guided Dataset Distillation with Balanced Distribution

AI summary

Overview

• Research area: Computer vision — dataset distillation (compressing large training sets into small synthetic ones), specifically the distribution-matching (DM) family of methods.
• Technical level: Intermediate. The paper assumes familiarity with dataset distillation, bilevel optimization, maximum mean discrepancy (MMD), and standard evaluation protocols such as images per class (IPC).
• Scope: The paper proposes Trajectory Guided Dataset Distillation (TGDD), which replaces static, randomly-initialized feature matching with alignment against model snapshots taken along pretraining trajectories, plus a stage-wise distribution constraint, and evaluates it on ten datasets.

What This Paper Is About

Distribution-matching dataset distillation synthesizes a tiny dataset by forcing its feature distribution to look like the real dataset's feature distribution under some network. Almost all such methods use a randomly initialized network as the feature extractor, so they only capture early-stage or poorly separated representations and ignore how features evolve as training progresses. The paper's goal is to make the matching process follow the model's training trajectory instead of a single static snapshot, producing synthetic data that is both diverse and class-separable, and therefore trains better downstream models.

Key Contributions

1. Reframing distribution matching as a dynamic balance. The authors reformulate DM as a dynamic process that jointly balances feature alignment and a distribution constraint across the whole training trajectory, rather than a single static matching step.
2. The TGDD method. They introduce Trajectory Guided Dataset Distillation, which uses multiple pretrained expert trajectories to perform (a) stage-wise distribution matching with a feature extractor sampled from a trajectory, and (b) a stage-wise distribution constraint applied by a second snapshot sampled from an "expert region" around that extractor.
3. Efficiency without extra optimization overhead. Expert trajectories are trained only on the original dataset and can be pretrained and reused, so TGDD keeps the speed of DM methods while using far fewer stored trajectories than trajectory-matching methods (the paper notes MTT stores 200 trajectories, while TGDD reports results with 5 trajectories).
4. Broad empirical validation. Experiments across ten datasets (SVHN, CIFAR-10, CIFAR-100, Tiny ImageNet, and the ImageNette, ImageWoof, ImageFruit, ImageMeow, ImageSquawk and ImageYellow subsets) report state-of-the-art results, including a 5.0% accuracy gain on high-resolution benchmarks.

Main Findings

• CIFAR-10 gains over DM: On CIFAR-10, TGDD reaches 71.9 ± 0.3 at IPC-10 and 76.5 ± 0.2 at IPC-50, versus 48.9 ± 0.6 and 63.0 ± 0.4 for the classical DM method. The paper states this is a 23% and 13.5% improvement respectively, and that it also outperforms optimization-oriented methods such as DC and MTT.
• Tiny ImageNet: TGDD achieves 29.3 ± 0.3 at IPC-10 and 30.9 ± 0.4 at IPC-50, which the authors describe as higher than current state-of-the-art methods.
• High-resolution ImageNet subsets: The abstract reports a 5.0% accuracy gain on high-resolution benchmarks; the text gives the concrete case of a 5% gain on ImageFruit at IPC-10 (34.8 ± 0.5 for TGDD versus 30.6 ± 0.8 for DANCE). TGDD also leads on ImageNette at IPC-10 (82 ± 0.5 versus 80.2 ± 0.7 for DANCE) and on ImageMeow at IPC-10 (41.4 ± 0.5 versus 39.4 ± 0.8 for DANCE).
• Cross-architecture generalization: Synthetic datasets distilled with a 3-layer ConvNet transfer to ResNet-10 and DenseNet-121. At IPC-50, TGDD scores 76.5 ± 0.2 (ConvNet-3), 74.9 ± 0.4 (ResNet-10) and 74.3 ± 0.2 (DenseNet-121), compared with 76.1 ± 0.1, 68.0 ± 0.1 and 64.8 ± 0.3 for DANCE.
• Both components matter: In the ablation on CIFAR-10/CIFAR-100, using only multi-formation augmentation gives 63.2/40.5 at IPC-10; adding stage-wise MMD gives 65.8/47.0; adding the distribution constraint gives 71.9/51.3. At IPC-50 the progression is 69.5 → 75.2 → 76.5 on CIFAR-10.
• Feature alignment across training stages: Figure 4(a) reports that TGDD consistently achieves better feature alignment at all training stages and configurations than the alternative, while using randomly initialized extractors instead (Figure 5a) produces overly concentrated feature distributions and hurts diversity.
• Class separability: Figure 4(b) shows TGDD's synthetic data yields higher classification accuracy across all pretrained models and stages, indicating reduced class overlap; replacing expert models with randomly initialized ones causes dispersed, entangled features (Figure 5b).
• Higher information density: Using the criterion that neurons with gradients exceeding layer-wise mean values count as activated, Figure 4(c) shows TGDD activates more neurons per image at equivalent total activation levels.
• Hyperparameter robustness: Performance fluctuates within 2.8% as the regularization coefficient α ranges from 0.5 to 10.5. The expert-region design performs consistently across region lengths L and clearly beats the single-expert setting; increasing the number or length of expert trajectories improves performance only slightly before saturating.
• Continual learning: With 20 images per class, step sizes of 5 and 10, and a 3-layer ConvNet, TGDD outperforms Herding, DSA, DM, DSDM and M3D (Figure 7).
• Qualitative quality: Visualizations (Figure 8) of CIFAR-10 and CIFAR-100 synthetic images show better structural detail retention, less noise, and more distinctive category-discriminative features than DM, M3D and DSDM.

Methodology in Plain English

1. Pretrain expert trajectories. Before distillation, the authors train N randomly initialized networks for M epochs each on the original dataset, saving a snapshot at every epoch to form trajectories of the form P = {p_i,j | 0 ≤ i ≤ N, 0 ≤ j ≤ M}. In the experiments these are 5 trajectories with 60 epochs for SVHN, CIFAR-10, CIFAR-100 and Tiny ImageNet, and 80 epochs for the ImageNet subsets.
2. Match distributions stage by stage. During distillation, each iteration samples a trajectory and then one snapshot from it, and uses that snapshot as the feature extractor. The synthetic and original batches are passed through it, and an MMD loss (L_MMD) aligns their per-class feature means. Because the snapshot changes each iteration, the synthetic data is matched against many different points along the learning trajectory instead of one random network.
3. Add an expert-region constraint. MMD constrains only the mean of the feature distribution, which the authors argue is too weak and leads to scattered synthetic features. So for each chosen extractor, they define an expert region of L consecutive snapshots starting from it, P_er = {p_i,j, ..., p_i,j+L-1}, sample one snapshot from that region, and apply a classification loss (L_SDC) to the synthetic images. This acts like an ensemble constraint with no extra training cost.
4. Combine and optimize. The total loss is L_overall = L_MMD + α·L_SDC, and only the synthetic dataset is updated, with D^S = D^S − η ∇L_overall. Settings used: α = 2.5 for 1 and 10 images per class and 0.5 for 50; L = 7 everywhere; distillation learning rate 0.1 for ImageNet subsets and 0.01 otherwise, scaled by IPC; SGD with 0.01 learning rate, 0.9 momentum, 0.0005 weight decay for training; multi-formation parameterization with factor ρ = 3 for ImageNet subsets and 2 otherwise; differential augmentation including color, crop, cutout, flip, scale and rotate. Each experiment was repeated 5 times.
5. Evaluate. Models are 3-layer ConvNets with 128 3×3 kernels, instance normalization, ReLU and 3×3 average pooling with stride 2 for SVHN/CIFAR-10/CIFAR-100, a 4-layer ConvNet for Tiny ImageNet, and a 5-layer ConvNet for the ImageNet subsets.

Why This Matters

Impact on research. The paper targets the central weakness of DM-based distillation — that it ignores how representations evolve — and shows that reusing cheaply pretrained trajectories fixes much of the gap with expensive optimization-oriented methods. This shifts attention from "which static encoder to use" toward "which points along training to match," and the cross-architecture results (training on a 3-layer ConvNet, testing on ResNet-10 and DenseNet-121) suggest the resulting synthetic data is not overfit to one architecture.

Real-world applications (the paper lists dataset distillation as benefiting these tasks):

• Neural architecture search, where many candidate models must be trained quickly.
• Federated learning, where communication and client-side storage are constrained.
• Continual learning, where compact replay sets must preserve earlier task performance — demonstrated here on CIFAR-100.
• Privacy protection, where only synthetic data rather than the original samples is shared or stored.

Industry relevance. The paper positions efficiency as a first-class result: Figure 1 compares accuracy, distillation time and GPU memory on CIFAR-10 under different IPCs with pretraining time included (5 trajectories for TGDD versus 100 for FTD). Because expert trajectories depend only on the original dataset, they can be pretrained once and reused across IPC settings, which makes the method attractive for pipelines where a distillation job is run repeatedly under changing resource budgets. The article notes the code is available at https://github.com/FlyFinley/TGDD.

Future Directions

• Other modalities and tasks. The conclusion states that future work will explore extending dataset distillation to other modalities and tasks beyond image classification.
• The pretraining cost trade-off. TGDD relies on pretrained expert trajectories; the paper does not report how the number or length of trajectories changes wall-clock time in detail, and Figures 6(c) and 6(d) show gains saturate as trajectories are added, leaving the optimal pretraining budget as an open question.
• Scaling to very large source datasets. The introduction frames the problem around data growth from ImageNet's tens of millions of images to LAION-5B's billions of image–text pairs; whether trajectory-based matching helps at that scale is not established here.
• Choosing the regularization strength automatically. Performance moves within 2.8% as α varies from 0.5 to 10.5, and the authors fix α by IPC (2.5 for 1 and 10 images per class, 0.5 for 50). A principled way to set α and the expert-region length L without per-setting tuning is not reported.

Target Audience

Researchers and practitioners working on dataset distillation, coreset selection, and data-efficient training, particularly those already familiar with DM-based methods and looking for a low-overhead way to improve them. It is also relevant to engineers who need to shrink training sets for storage- or compute-constrained deployment, and to readers interested in how training dynamics (rather than a single static representation) can be exploited in data synthesis. Readers new to the field will need to read up on MMD, bilevel optimization, and IPC-based evaluation protocols first.

Authors’ abstract

Dataset distillation compresses large datasets into compact synthetic ones to reduce storage and computational costs. Among various approaches, distribution matching (DM)-based methods have attracted attention for their high efficiency. However, they often overlook the evolution of feature representations during training, which limits the expressiveness of synthetic data and weakens downstream performance. To address this issue, we propose Trajectory Guided Dataset Distillation (TGDD), which reformulates distribution matching as a dynamic alignment process along the model's training trajectory. At each training stage, TGDD captures evolving semantics by aligning the feature distribution between the synthetic and original dataset. Meanwhile, it introduces a distribution constraint regularization to reduce class overlap. This design helps synthetic data preserve both semantic diversity and representativeness, improving performance in downstream tasks. Without additional optimization overhead, TGDD achieves a favorable balance between performance and efficiency. Experiments on ten datasets demonstrate that TGDD achieves state-of-the-art performance, notably a 5.0% accuracy gain on high-resolution benchmarks.