Data Distillation: A Survey

Meta-model matching-based data distillation approaches fundamentally optimize for the transferability of models trained on the data summary when generalized to the original dataset:

where intuitively, the inner-loop trains a representative learning algorithm on the data summary until convergence, and the outer-loop subsequently optimizes the data summary for the transferability of the optimized learning algorithm to the original dataset. Besides common assumptions mentioned earlier, the key simplifying assumption for this family of methods is that a perfect classifier exists and can be estimated on $\mathcal{D}$, i.e., $\exists~{}\theta^{\mathcal{D}}$ s.t. $l(\Phi_{\theta^{\mathcal{D}}}(x),y)=0,~{}\forall x\sim\mathcal{X},y\sim\mathcal{Y}$. Plugging the second assumption along with the iid assumption of $\mathcal{D}$ in Equation 2 directly translates to Equation 3. Despite the assumption, Equation 3 is highly expensive both in terms of computation time and memory, due to which, methods from this family typically resort to making further assumptions.

Wang et al. (2018) (DD) originally proposed the task of data distillation, and used the meta-model matching framework for optimization. DD makes the optimization in Equation 3 tractable by performing (1) local optimization à la stochastic gradient descent (SGD) in the inner-loop, and (2) outer-loop optimization using Truncated Back-Propagation Through Time (TBPTT), i.e., unrolling only a limited number of inner-loop optimization steps. Formally, the modified optimization objective for DD is as follows:

where $\mathbf{P}_{\theta}$ is a parameter initialization distribution of choice, $T$ accounts for the truncation in TBPTT, and $\eta$ is a tunable learning rate. We also elucidate DD’s control-flow in Algorithm 1 for reference.

Notably, TBPTT has been associated with drawbacks such as (1) computationally expensive inner-loop unrolling; (2) bias involved with truncated unrolling (Wu et al., 2018); and (3) poorly conditioned loss landscapes, particularly with long unrolls (Metz et al., 2019). Consequently, the TBPTT framework was empirically shown to be ineffective for data distillation (Zhao et al., 2021). However, recent work (Deng & Russakovsky, 2022) claims that using momentum-based optimizers and longer inner-loop unrolling can greatly improve performance. We delay a deeper discussion of this work to Section 2.5 for clarity.

Analogously, a separate line of work focuses on using Neural Tangent Kernel (NTK) (Jacot et al., 2018) based algorithms to solve the inner-loop in closed form. As a brief side note, the infinite-width correspondence states that performing Kernelized Ridge Regression (KRR) using the NTK of a given neural network, is equivalent to training the same $\infty$-width neural network with L2 reconstruction loss for $\infty$ SGD-steps. These “$\infty$-width” neural networks have been shown to perform reasonably compared to their finite-width counterparts, while also being solved in closed-form (see Lee et al. (2020) for a detailed analysis on finite vs. infinite neural networks for image classification). KIP uses the NTK of a fully-connected neural network (Nguyen et al., 2021a), or a convolutional network (Nguyen et al., 2021b) in the inner-loop of Equation 3 for efficient data distillation. More formally, given the NTK $\mathcal{K}:\mathcal{X}\times\mathcal{X}\mapsto\mathbb{R}$ of a neural network architecture, KIP optimizes the following objective:

where $\mathbf{K}_{AB}\in\mathbb{R}^{|A|\times|B|}$ represents the gramian matrix of two sets $A$ and $B$, and whose $(i,j)^{\text{th}}$ element is defined by $\mathcal{K}(A_{i},B_{j})$. Although KIP doesn’t impose any additional simplifications to the meta-model matching framework, it has an $\mathcal{O}(|\mathcal{D}|\cdot n\cdot\dim(\mathcal{X}))$ time and memory complexity, limiting its scalability. Subsequently, RFAD (Loo et al., 2022) proposes using (1) the light-weight Empirical Neural Network Gaussian Process (NNGP) kernel (Neal, 2012) instead of the NTK; and (2) a classification loss (e.g., NLL) instead of the L2-reconstruction loss for the outer-loop to get $\mathcal{O}(n)$ time complexity while also having better performance. On a similar note, FRePO (Zhou et al., 2022b) decouples the feature extractor and a linear classifier in $\Phi$, and alternatively optimizes (1) the data summary along with the classifier, and (2) the feature extractor. To be precise, let $f_{\theta}:\mathcal{X}\mapsto\mathcal{X}^{\prime}$ be the feature extractor, $g_{\psi}:\mathcal{X}^{\prime}\mapsto\mathcal{Y}$ be the linear classifier, s.t. $\Phi(x)\equiv g_{\psi}(f_{\theta}(x))~{}\forall x\in\mathcal{X}$; the optimization objective for FRePO can be written as:

where $T$ represents the number of inner-loop update steps for the feature extractor $f_{\theta}$. Notably, (1) a wide architecture for $f_{\theta}$ is crucial for distillation quality in FRePO; and (2) despite the bilevel optimization, FRePO is shown to be more scalable compared to KIP (Equation 5), while also being more generalizable.

Gradient matching based data distillation, at a high level, performs one-step distance matching on (1) the network trained on the target dataset ($\mathcal{D}$) vs. (2) the same network trained on the data summary ($\mathcal{D}_{\mathsf{syn}}$). In contrast to the meta-model matching framework, such an approach circumvents the unrolling of the inner-loop, thereby making the overall optimization much more efficient. First proposed by Zhao et al. (2021) (DC), data summaries optimized by gradient-matching significantly outperformed data pruning methodologies, as well as TBPTT-based data distillation proposed by Wang et al. (2018). Formally, given a learning algorithm $\Phi$, DC solves the following optimization objective:

where $T$ accounts for model similarity $T$-steps in the future, and $\mathbf{D}:\mathbb{R}^{|\theta|}\times\mathbb{R}^{|\theta|}\mapsto\mathbb{R}$ is a distance metric of choice (typically cosine distance). In addition to assumptions imposed by the meta-model matching framework (Section 2.1), gradient-matching assumes (1) inner-loop optimization of only $T$ steps; (2) local smoothness: two sets of model parameters close to each other (given a distance metric) imply model similarity; and (3) first-order approximation of $\theta_{t}^{\mathcal{D}}$: instead of exactly computing the training trajectory of optimizing $\theta_{0}$ on $\mathcal{D}$ (say $\theta_{t}^{\mathcal{D}}$); perform first-order approximation on the optimization trajectory of $\theta_{0}$ on the much smaller $\mathcal{D}_{\mathsf{syn}}$ (say $\theta_{t}^{\mathcal{D}_{\mathsf{syn}}}$), i.e., approximate $\theta_{t}^{\mathcal{D}}$ as a single gradient-descent update on $\theta_{t-1}^{\mathcal{D}_{\mathsf{syn}}}$ using $\mathcal{D}$ rather than $\theta_{t-1}^{\mathcal{D}}$ (Figure 3).

Subsequently, numerous other approaches have been built atop this framework with subtle variations. DSA (Zhao & Bilen, 2021) improves over DC by performing the same image-augmentations (e.g., crop, rotate, jitter, etc.) on both $\mathcal{D}$ and $\mathcal{D}_{\mathsf{syn}}$ while optimizing Equation 7. Since these augmentations are universal and are applicable across data distillation frameworks, augmentations performed by DSA have become a common part of all methods proposed henceforth, but we omit them for notational clarity. DCC (Lee et al., 2022b) further modifies the gradient-matching objective to incorporate class contrastive signals inside each gradient-matching step and is shown to improve stability as well as performance. With $\theta_{t}$ evolving similarly as in Equation 7, the modified optimization objective for DCC can be written as:

Most recently, Kim et al. (2022) (IDC) extend the gradient matching framework by: (1) multi-formation: to synthesize a higher amount of data within the same memory budget, store the data summary (e.g., images) in a lower resolution to remove spatial redundancies, and upsample (using e.g., bilinear, FSRCNN (Dong et al., 2016)) to the original scale while usage; and (2) matching gradients of the network’s training trajectory over the full dataset $\mathcal{D}$ rather than the data summary $\mathcal{D}_{\mathsf{syn}}$. To be specific, given a $k\times$ upscaling function $f:\mathbb{R}^{d\times d}\mapsto\mathbb{R}^{kd\times kd}$, the modified optimization objective for IDC can be formalized as:

Kim et al. (2022) further hypothesize that training models on $\mathcal{D}_{\mathsf{syn}}$ instead of $\mathcal{D}$ in the inner-loop has two major drawbacks: (1) strong coupling of the inner- and outer-loop resulting in a chicken-egg problem (McLachlan & Krishnan, 2007); and (2) vanishing network gradients due to the small size of $\mathcal{D}_{\mathsf{syn}}$, leading to an improper outer-loop optimization for gradient-matching based techniques.

Cazenavette et al. (2022) proposed MTT which aims to match the training trajectories of models trained on $\mathcal{D}$ vs. $\mathcal{D}_{\mathsf{syn}}$. More specifically, let $\{\theta_{t}^{\mathcal{D}}\}_{t=0}^{T}$ represent the training trajectory of training $\Phi_{\theta}$ on $\mathcal{D}$; trajectory matching algorithms aim to solve the following optimization:

where $\mathbf{D}:\mathbb{R}^{|\theta|}\times\mathbb{R}^{|\theta|}\mapsto\mathbb{R}$ is a distance metric of choice (typically L2 distance). Such an optimization can intuitively be seen as optimizing for similar quality models trained with $N$ SGD steps on $\mathcal{D}_{\mathsf{syn}}$, compared to $M\gg N$ steps on $\mathcal{D}$, thereby invoking long-horizon trajectory matching. Notably, calculating the gradient of Equation 10 w.r.t. $\mathcal{D}_{\mathsf{syn}}$ encompasses gradient unrolling through $N$-timesteps, thereby limiting the scalability of MTT. On the other hand, since the trajectory of training $\Phi_{\theta}$ on $\mathcal{D}$, i.e., $\{\theta_{t}^{\mathcal{D}}\}_{t=0}^{T}$ is independent of the optimization of $\mathcal{D}_{\mathsf{syn}}$, it can be pre-computed for various $\theta_{0}\sim\mathbf{P}_{\theta}$ initializations and directly substituted. Similar to gradient matching methods (Section 2.2), the trajectory matching framework also optimizes the first-order distance between parameters, thereby inheriting the local smoothness assumption. As a scalable alternative, Cui et al. (2022b) proposed TESLA, which re-parameterizes the parameter-matching loss of MTT in Equation 10 (specifically when $\mathbf{D}$ is set as the L2 distance), using linear algebraic manipulations to make the bilevel optimization’s memory complexity independent of $N$. Furthermore, TESLA uses learnable soft-labels ($\mathbf{Y}_{\mathsf{syn}}$) during the optimization for an increased compression efficiency.

Even though the aforementioned gradient-matching or trajectory-matching based data distillation techniques have been empirically shown to synthesize high-quality data summaries, the underlying bilevel optimization, however, is oftentimes a computationally expensive procedure. To this end, distribution-matching techniques solve a correlated proxy task via a single-level optimization, leading to a vastly improved scalability. More specifically, instead of matching the quality of models on $\mathcal{D}$ vs. $\mathcal{D}_{\mathsf{syn}}$, distribution-matching techniques directly match the distribution of $\mathcal{D}$ vs. $\mathcal{D}_{\mathsf{syn}}$. The key assumption for this family of methods is that two datasets that are similar according to a particular distribution divergence metric, lead to similarly trained models.

First proposed by Zhao & Bilen (2023), DM uses (1) numerous parametric encoders to cast high-dimensional data into respective low-dimensional latent spaces; and (2) an approximation of the Maximum Mean Discrepancy to compute the distribution mismatch between $\mathcal{D}$ and $\mathcal{D}_{\mathsf{syn}}$ in each of the latent spaces. More precisely, given a set of $k$ encoders $\mathcal{E}\triangleq\{\psi_{i}:\mathcal{X}\mapsto\mathcal{X}_{i}\}_{i=1}^{k}$, the optimization objective can be written as:

DM uses a set of randomly initialized neural networks (with the same architecture) to instantiate $\mathcal{E}$. They observe similar performance when instantiated with more meaningful, task-optimized neural networks, despite it being much less efficient. CAFE (Wang et al., 2022) further refines the distribution-matching idea by: (1) solving a bilevel optimization problem for jointly optimizing a single encoder ($\Phi$) and the data summary, rather than using a pre-determined set of encoders ($\mathcal{E}$); and (2) assuming a neural network encoder ($\Phi$), match the latent representations obtained at all intermediate layers of the encoder instead of only the last layer. Formally, given a $(L+1)$-layer neural network $\Phi_{\theta}:\mathcal{X}\mapsto\mathcal{Y}$ where $\Phi^{l}_{\theta}$ represents $\Phi$’s output at the $l^{\text{th}}$ layer, the optimization problem for CAFE can be specified as:

where $\hat{p}(\cdot|\cdot,\theta)$ intuitively represents the nearest centroid classifier on $\mathcal{D}_{\mathsf{syn}}$ using the latent representations obtained by last layer of $\Phi_{\theta}$. Analogously, IT-GAN (Zhao & Bilen, 2022) also uses the distribution-matching framework in Equation 11 to generate data that is informative for model training, in contrast to the traditional GAN (Goodfellow et al., 2014) which focuses on generating realistic data.

All of the aforementioned data distillation frameworks intrinsically maintain the synthesized data summary as a large set of free parameters, which are in turn optimized. Arguably, such a setup prohibits knowledge sharing between synthesized data points (parameters), which might introduce data redundancy. On the other hand, factorization-based data distillation techniques parameterize the data summary using two separate components: (1) bases: a set of mutually independent base vectors; and (2) hallucinators: a mapping from the bases’ vector space to the joint data- and label-space. In turn, both the bases and hallucinators are optimized for the task of data distillation.

Formally, let $\mathcal{B}\triangleq\{b_{i}\in\mathbb{B}\}_{i=1}^{|\mathcal{B}|}$ be the set of bases, and $\mathcal{H}\triangleq\{h_{i}:\mathbb{B}\mapsto\mathcal{X}\times\mathcal{Y}\}_{i=1}^{|\mathcal{H}|}$ be the set of hallucinators, then the data summary is parameterized as $\mathcal{D}_{\mathsf{syn}}\triangleq\{{h(b)}\}_{b\sim\mathcal{B},~{}h\sim\mathcal{H}}$. Even though such a two-pronged approach seems similar to generative modeling of data, note that unlike classic generative models, (1) the input space consists only of a fixed and optimized set of latent codes and isn’t meant to take any other inputs; and (2) given a specific $\mathcal{B}$ and $\mathcal{H}$, we can generate at most $|\mathcal{B}|\cdot|\mathcal{H}|$ sized data summaries. Notably, such a hallucinator-bases data parameterization can be optimized using any of the aforementioned data optimization frameworks (Sections 2.1, 2.2, 2.4 and 2.3)

This framework was concurrently proposed by Deng & Russakovsky (2022) (we take the liberty to term their unnamed model as “Lin-ear Ba-ses”) and Liu et al. (2022c) (HaBa). LinBa modifies the general hallucinator-bases framework by assuming (1) the bases’ vector space ($\mathbb{B}$) to be the same as the task input space ($\mathcal{X}$); and (2) the hallucinator to be linear and additionally conditioned on a given predictand. More specifically, the data parameterization can be formalized as follows:

where for the sake of notational simplicity, we assume $y\in\mathbb{R}^{|\mathcal{C}|}$ represents the one-hot vector of the label for which we want to generate data, and the maximum amount of data that can be synthesized $n\leq|\mathcal{C}|\cdot|\mathcal{H}|$. Since the data generation (Equation 13) is end-to-end differentiable, both $\mathbf{B}$ and $\mathcal{H}$ are jointly optimized using the TBPTT framework discussed in Section 2.1, albeit with some crucial modifications for vastly improved performance: (1) using momentum-based optimizers instead of vanilla SGD in the inner-loop; and (2) longer unrolling ($\geq 100$ steps) of the inner-loop during TBPTT. Liu et al. (2022c) (HaBa) relax the linear and predictand-conditional hallucinator assumption of LinBa, equating to the following data parameterization:

where $\mathcal{B}$ and $\mathcal{H}$ are optimized using the trajectory matching framework (Section 2.3) with an additional contrastive constraint to promote diversity in $\mathcal{D}_{\mathsf{syn}}$ (cf. Liu et al. (2022c), Equation (6)). Following this setup, HaBa can generate at most $|\mathcal{B}|\cdot|\mathcal{H}|$ sized data summaries. Furthermore, one striking difference between HaBa (Equation 14) and LinBa (Equation 13) is that to generate each data point, LinBa uses a linear combination of all the bases, whereas HaBa generates a data point using a single base vector.