PaperlayerView exact revision on arXiv

Source-linked AI summary

Remember the Past: Distilling Datasets into Addressable Memories for Neural Networks

Zhiwei Deng, Olga Russakovsky

arXiv:2206.02916v2cs.LGcs.AIcs.CV

TL;DR

Dataset distillation seeks compact training data that preserves model performance, while class-specific representations can scale with class count and miss shared information. The paper distills shared bases into addressable memories, reporting strong results across distillation and continual-learning benchmarks.

  • Problem

    Dataset distillation seeks to compress training data while preserving model generalization, but class-specific representations grow with class count and ignore shared information across classes.

    Source: §3 · p. 3

  • Method

    The method learns common memory bases shared across classes and combines them with addressing functions to generate synthetic training examples and support general queries.

    Source: §1 · p. 2, §4 · p. 3

  • Results

    The method consistently outperforms prior state of the art across six dataset-distillation benchmarks and achieves state of the art on four continual-learning datasets, including 66.4% accuracy on CIFAR10 with 1 image per class, 16.5% above KIP.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The shared representation removes the need for compressed data size to grow linearly with class count and permits queries beyond categorical labels.

    Source: §4 · p. 3, §6 · p. 10

    The costly inner optimization loop may be time-consuming on larger models or datasets.

    Source: §6 · p. 10

Abstract

from arXiv · show

We propose an algorithm that compresses the critical information of a large dataset into compact addressable memories. These memories can then be recalled to quickly re-train a neural network and recover the performance (instead of storing and re-training on the full original dataset). Building upon the dataset distillation framework, we make a key observation that a shared common representation allows for more efficient and effective distillation. Concretely, we learn a set of bases (aka ``memories'') which are shared between classes and combined through learned flexible addressing functions to generate a diverse set of training examples. This leads to several benefits: 1) the size of compressed data does not necessarily grow linearly with the number of classes; 2) an overall higher compression rate with more effective distillation is achieved; and 3) more generalized queries are allowed beyond recalling the original classes. We demonstrate state-of-the-art results on the dataset distillation task across six benchmarks, including up to 16.5% and 9.7% in retained accuracy improvement when distilling CIFAR10 and CIFAR100 respectively. We then leverage our framework to perform continual learning, achieving state-of-the-art results on four benchmarks, with 23.2% accuracy improvement on MANY. The code is released on our project webpage https://github.com/princetonvisualai/RememberThePast-DatasetDistillation.

1 Introduction

The paper addresses the limits of class-specific distilled examples by proposing shared addressable memories, then evaluates them for dataset distillation and continual learning. It reports strong benchmark results and identifies training choices that improve the method.

1 Introduction

  • Dataset distillation learns compressed examples that summarize training-relevant information so models can be retrained to recover skills, unlike standard reconstruction-based compression.The paper motivates this with deep networks' tendency to forget past knowledge when learning new tasks.Source: §1 · p. 1
  • Class-specific learnable examples can make memory size grow linearly with class count, potentially creating redundancies and limiting generalization to continuous label spaces.The paper notes that continuous label spaces can contain infinitely many values.Source: §1 · pp. 1–2
  • Shared bases combined through learned addressing functions form compact memories that generate training examples, exploiting information shared between classes.The formulation can improve compression and performance while allowing queries beyond categorical labels.Source: §1 · p. 2
  • Back-propagation through time trains the memories and addressing functions; momentum and long inner-loop unrolls are critical factors for performance.This differs from common bi-level optimization use on the task and enables strong performance over single-step gradient matching, even with simple data representations.Source: §1 · p. 2
  • 66.4% accuracy on CIFAR10 with storage of 1 image per class improves over KIP by 16.5%, and the method outperforms prior state of the art across six distillation benchmarks.The paper also reports state-of-the-art continual-learning results on four datasets, including a 23.2% retained-accuracy improvement on MANY.Source: §1 · p. 2

2 Related works

Related work surveys dataset-distillation criteria and continual-learning approaches to forgetting. It situates the paper among methods that compress or preserve information through optimization, matching, kernels, buffers, and stored bases.

2 Related works

  • Dataset-distillation methods prioritize information using final-state generalization loss, gradient or score matching, or kernel inducing points.The cited kernel methods connect to Gaussian processes but can incur large computation costs.Source: §2 · p. 2
  • Dataset-distillation methods are connected to pragmatic compression approaches that compress or match distributions based on a decision process.For dataset distillation, the decision process is the gradient-descent search process.Source: §2 · p. 2
  • Continual-learning methods address catastrophic forgetting and the stability-plasticity dilemma through regularization, replay, or dynamic models.Memory buffers have often stored randomly selected real samples using different strategies.Source: §2 · p. 2
  • Recent continual-learning work also stores random bases for online settings or SVD bases for offline settings.These approaches are presented alongside memory buffers as ways of using memory in continual learning.Source: §2 · p. 2

3 Background: dataset distillation

Dataset distillation compresses a large training dataset into a small learned dataset intended to preserve model generalization. Existing class-balanced learnable-example representations can scale with class count, ignore shared information, and fail to generalize to new classes or tasks.

3 Background: dataset distillation

  • Dataset distillation learns a small synthetic dataset intended to preserve the generalization ability of a model trained on the original training data.The setup compares model performance on a test dataset using a metric such as accuracy.Source: §3 · pp. 2–3
  • Existing synthetic representations use learnable data samples distributed equally across classes, but their sample count grows linearly with the number of classes.This limits applicability when class counts are large or undefined, including language or other continuous labels.Source: §3 · p. 3
  • Class-specific representations ignore potentially shared information and do not learn a common representation, reducing compactness and limiting generalization to new classes or tasks.Source: §3 · p. 3

4 Model

The model recasts dataset distillation as shared memories accessed by learned functions, then trains those memories and functions through back-propagation through time. The framework supports flexible queries and reports that its BPTT improvements outperform single-step gradient-matching methods.

4 Model

  • The improved BPTT framework reports drastic performance gains and outperforms single-step gradient-matching methods.The overview contrasts this result with the common observation that gradient matching outperforms BPTT on dataset distillation.Source: §4 · p. 3

4.1 Dataset Distillation as memory addressing

  • Shared bases and learned addressing replace separate class-specific synthetic examples, so memory size need not grow linearly with class count.The bases represent shared dataset components, while the storage budget includes both memories and addressing-function parameters.Source: §4.1 · pp. 3–4
  • For each query, parameterized matrices produce coefficients that linearly recombine the shared bases into synthetic examples paired with that query.The query can be a vectorized label representation, including a one-hot categorical label.Source: §4.1 · p. 4
  • Queries can also be continuous or sampled from a feature-extractor-defined distribution, rather than being limited to enumerated discrete labels.The paper gives language and audio as examples of possible query forms.Source: §4.1 · p. 4

Algorithm 1

  • Algorithm 1 recalls synthetic data for sampled labels, trains a randomly initialized model on synthetic minibatches, and back-propagates the generalization-loss gradient through the inner loop to update the memories and addressing function.The inner loop uses momentum-based model updates; the outer loss is computed on a minibatch from the training dataset.Source: Algorithm 1 · p. 5, §4.2 · p. 5

4.2 Learning framework: back-propagation through time

  • Momentum is critical even for short inner loops, while sufficiently long unrolled trajectories—such as 200 steps—with momentum can produce more effective gradients for learning the compressed representation.Earlier uses of BPTT on this task used relatively short trajectories, such as 10–30 steps.Source: §4.2 · p. 6

5 Experiments

Across six dataset-distillation benchmarks, the experiments report state-of-the-art performance; a simple compress-then-recall approach also outperforms prior methods on four continual-learning benchmarks. The compressed representation further supports classifiers for new task combinations and image-feature queries.

5.1 Dataset Distillation

  • 87.3% on SVHN and 66.4% on CIFAR10 with 1 image per class outperform prior methods by 30% and 16.5%, respectively, under the same storage budget.The paper reports that these results also beat prior methods using 10 images per class.Source: Table 1 · p. 7, §5.1 · p. 7
  • In CIFAR100, tree classes share addressing bases strongly, while lawn mower and rocket are distinct; similar sharing patterns appear in CIFAR10.Higher cosine similarity indicates that classes use similar memory-basis components to produce synthetic images.Source: §5.1 · p. 7, Figure 3 · p. 8
  • 18.2% on CIFAR10 and 42.3% on SVHN are the recovered-accuracy increases reported for longer inner-loop trajectories versus 1-step cases.Momentum also consistently boosts performance with both 10-step and 100-step inner loops; the examples given are 7.0% and 9.2% on CIFAR10.Source: §5.1 · p. 8, Figure 4 · p. 8

5.2 Continual learning

  • 23% boost on MNIST MANY, from 50.81% to 74.07%, accompanies results that outperform prior methods on four continual-learning benchmarks.Table 3 summarizes comparisons with previous methods under the same memory sizes.Source: Table 3 · p. 9, §5.2 · p. 9
  • 87.3±0.92 on Permutated MNIST and 88.3±0.58 on Rotated MNIST exceed KCL, Stable SGD, and the reported multitask upper bounds in the 60,000-sample-per-task setting.The corresponding KCL scores are 85.5±0.78 and 81.8±0.60; Stable SGD scores are 80.1±0.51 and 70.8±0.78; multitask upper bounds are 86.5±0.21 and 87.3±0.47.Source: §5.2 · p. 9

5.3.1 Extrapolating between tasks

  • 72.53%±8.74 on 2-way and 46.54%±6.42 on 5-way classification show that recalled data can train classifiers for CIFAR100 class combinations not seen together during training.The task combinations are formed by selecting one class from each of randomly chosen training tasks; the storage budget is 1 image per class.Source: §5.3.1 · p. 10

5.3.2 Dataset Distillation extension – recall the past with images

  • The framework performs image addressing from visual observations, and the resulting model outperforms two strong baselines on constructing new classifiers on CIFAR100.The setup considers 1-shot and 5-shot observations; image features provide query vectors for recalling memories.Source: §5.3.2 · p. 10

6 Conclusion and limitations

The paper concludes that compact addressable memories support dataset distillation and continual learning, while noting computational and societal limitations.

Conclusion

  • The framework removes linear growth constraints on compressed data size, supports queries beyond categorical labels, and achieves strong retraining performance.It also leads to state-of-the-art continual-learning results on four datasets.Source: §6 · p. 10

Limitations

  • The full model's costly inner optimization loop may be time-consuming on larger models or datasets.The paper suggests combining the memory formulation with a different learning framework as a possible remedy.Source: §6 · p. 10
  • Distilled datasets may omit diversity in the original distribution, potentially harming retrained classifier performance on minority populations.The authors state that improved retrained accuracy arguably takes a step toward mitigating this concern.Source: §6 · p. 10

A Experiment setups

This appendix section introduces detailed experimental setups for the tasks discussed in the main paper.

A Experiment setups

  • The appendix provides detailed setups for the main-paper tasks, covering datasets, architectures, and implementation details.The section is an overview of the experimental setup material.Source: A Experiment setups · p. 14
  • The setup details apply to all tasks discussed in the main paper.Source: A Experiment setups · p. 14
  • The section explains experimental configurations rather than reporting results.Source: A Experiment setups · p. 14

A.1 Dataset Distillation

The dataset-distillation experiments use six standard benchmarks and describe their architectures and implementation settings.

A.1 Dataset Distillation

  • The experiments cover six benchmarks: MNIST, FashionMNIST, SVHN, CIFAR10, CIFAR100, and TinyImageNet.The benchmarks span digit, clothing, street-number, and object-image classification, with CIFAR100 and TinyImageNet covering 100 and 200 classes, respectively.Source: §A.1 · p. 14
  • The main model is a three-layer ConvNet, with additional evaluation on a modified ResNet-12.The ResNet-12 replaces BatchNorm with InstanceNorm and removes the final average-pooling layer; the paper says full final-layer spatial information is important for distillation.Source: §A.1 · p. 14
  • Experiments use SGD with learning rate 0.1 and momentum 0.5, and models train for 50,000 iterations.CIFAR10, CIFAR100, and SVHN use ZCA preprocessing; CIFAR10 and CIFAR100 also use rotation and flip augmentation.Source: §A.1 · p. 14

A.2 Continual learning

The appendix specifies continual-learning datasets and implementations, then describes CIFAR100 setups for task extrapolation and recall with images.

A.2 Continual learning

  • Continual-learning experiments use MLPs for MNIST-based datasets and a 160-hidden-dimension ConvNet following La-MAML.They use SGD with learning rate 0.1 and momentum 0.5; testing retraining follows the inner-loop setup.Source: §A.2 · p. 14
  • The continual-learning evaluation uses six datasets, with benchmark details summarized in Table 5.The listed settings include Rotations, Permutations, MANY, Incremental CIFAR100, and larger-sample Rotations and Permutations variants.Source: §A.2 · p. 14, Table 5 · p. 15

A.3 New classifier synthesis

  • Task extrapolation and image recall are evaluated on CIFAR100 using distinct train-test task constructions.Extrapolation trains on 20 five-way tasks and tests on two- or five-way tasks whose class pairs were not used together in training; image recall trains on all classes and tests on 20 five-way tasks using test-set image shots.Source: §A.3 · p. 15
  • For image recall, the model uses 64 bases and 16 addressing matrices, while the nearest-neighbor and classify-then-recall baselines are pretrained for 100 epochs.The visual observations used for recall come from the test set.Source: §A.3 · p. 15

B.1 Dataset Distillation

The section reports dataset-distillation comparisons and design choices, then describes memory allocation for continual learning and image-based addressing for new classifiers. It also discusses qualitative results for classifier synthesis.

B.1 Dataset Distillation

  • Back-propagation through time consistently outperforms single-step gradient matching, while downsampling reduces spatial redundancy and improves compression rate and recovery performance.The reported benchmark results are in Table 7.Source: §B.1 · p. 15, Table 7 · p. 16
  • The learned compressed representation is reported to be robust across ConvNet and ResNet-12 when trained on one architecture and tested on the other.The comparison uses 10 images per class on CIFAR10; Table 6 summarizes results across pixel/image storage budgets.Source: Table 6 · p. 15, §B.1 · p. 16
  • The number of bases is selected by choosing the validation-set setting with the highest performance, using a random 10% of training samples for validation.Figure 5 plots base count against retraining accuracy on the validation set.Source: §B.1 · p. 16, Figure 5 · p. 16
  • The momentum ablation compares BPTT without momentum, momentum used only in the forward pass, and the full model, on CIFAR10 with 200 inner optimization steps.The forward-only variant blocks momentum gradients in the backward pass except for gradients on current-step weights; Table 8 is the further momentum analysis.Source: §B.1 · p. 16, Table 8 · p. 16
  • Adam caused gradient instability at the same learning-rate magnitude as SGD with momentum, while smaller rates such as 1e-4 produced similar end results.The comparison concerns optimizing the synthetic data in the inner loop.Source: §B.1 · p. 17

B.2 Continual learning

  • Continual learning stores reservoir-sampled examples in the buffer, then replaces current-task real samples with compressed representations allocated 1/T of the buffer after each task.During task learning, the algorithm uses available buffer space for samples; at task index t, the stated available fraction is 1 −(t −1)/T.Source: §B.2 · p. 17

B.3 New classifier synthesis

  • For new-classifier synthesis, image features act as queries to address shared bases, and the reported comparison notes that direct feature queries can yield more robust recall than classify-then-recall.The feature extractor, bases, and image matrices are jointly trained; classify-then-recall can be less robust because test-image classification errors affect recall.Source: §B.3 · p. 17

C Visualization and analysis

The visual analyses show similarities in how classes combine bases and depict the learned bases and recalled synthetic images. Additional comparisons vary the number of generated images per class and the corresponding number of bases.

C Visualization and analysis

  • CIFAR100 classes within the same superclass often have similar basis-combination coefficients, while crab and tulips use very different coefficients.The coefficient cosine-similarity matrix orders classes by superclass, with every five classes sharing a superclass.Source: C Visualization and analysis · p. 17, Figure 6 · p. 18
  • The 64 learned CIFAR100 bases contain varied colors, shapes, and textures, and are combined using coefficients generated from queries and addressing matrices.Figure 7 visualizes the learned bases.Source: C Visualization and analysis · p. 17, Figure 7 · p. 19

D.1 Same amount of generated images

  • The equal-image-count visualization compares BPTT at 100 images per class with memory addressing at 10 images per class and 43 bases, producing approximately 99 recalled images.Neither method uses downsampling in this comparison; the synthetic images are shown in the accompanying figures.Source: §D.1 · pp. 17–18, Figure 8 · p. 19

D.2 Various image per class budgets

  • The image-budget visualizations compare 2, 10, and 50 images per class using 8, 43, and 215 bases, respectively.Figures 12–14 show recalled images for apple, aquarium fish, and baby across these settings; Figures 10–11 show bed and bottle.Source: §D.2 · p. 19, Figure 10 · p. 20, Figure 11 · p. 21
Loading 2206.02916v2…