Dataset Distillation Efficiently Encodes Low-Dimensional Representations from Gradient-Based Learning of Non-Linear Tasks

Dataset distillation, a training-aware data compression technique, has recently attracted increasing attention as an effective tool for mitigating costs of optimization and data storage. However, progress remains largely empirical. Mechanisms underlying the extraction of task-relevant information from the training process and the efficient encoding of such information into synthetic data points remain elusive. In this paper, we theoretically analyze practical algorithms of dataset distillation applied to the gradient-based training of two-layer neural networks with width $L$. By focusing on a non-linear task structure called multi-index model, we prove that the low-dimensional structure of the problem is efficiently encoded into the resulting distilled data. This dataset reproduces a model with high generalization ability for a required memory complexity of $\tilde \Theta(r^2d+L)$, where $d$ and $r$ are the input and intrinsic dimensions of the task. To the best of our knowledge, this is one of the first theoretical works that include a specific task structure, leverage its intrinsic dimensionality to quantify the compression rate and study dataset distillation implemented solely via gradient-based algorithms.

Training modern artificial intelligence (AI) often requires very large datasets, which can make learning expensive in time, storage, and energy. A promising idea called dataset distillation tries to replace a large dataset with a much smaller set of artificial examples that still teaches an AI model the same task. Although this method often works well in practice, researchers have not fully understood why such tiny datasets can contain enough useful information. In this work, we study this question mathematically for a common type of neural network and training procedure. We focus on tasks where the important information relies only on a few aspects of the raw data. We show that dataset distillation can efficiently capture this low-dimensional structure and store it as synthetic examples, so that only a few such examples are needed for the model to generalize well to new data. This helps explain when and why dataset distillation works, and may guide the design of cheaper, more efficient ways to train AI systems.