ATOM: Attention Mixer for Efficient Dataset Distillation

Abstract

Recent works in dataset distillation seek to minimize training expenses by generating a condensed synthetic dataset that encapsulates the information present in a larger real dataset. These approaches ultimately aim to attain test accuracy levels akin to those achieved by models trained on the entirety of the original dataset. Previous studies in feature and distribution matching have achieved significant results without incurring the costs of bi-level optimization in the distillation process. Despite their convincing efficiency, many of these methods suffer from marginal downstream performance improvements, limited distillation of contextual information, and subpar cross-architecture generalization. To address these challenges in dataset distillation, we propose the ATtentiOn Mixer (ATOM) module to efficiently distill large datasets using a mixture of channel and spatial-wise attention in the feature matching process. Spatial-wise attention helps guide the learning process based on consistent localization of classes in their respective images, allowing for distillation from a broader receptive field. Meanwhile, channel-wise attention captures the contextual information associated with the class itself, thus making the synthetic image more informative for training. By integrating both types of attention, our ATOM module demonstrates superior performance across various computer vision datasets, including CIFAR10/100 and TinyImagenet. Notably, our method significantly improves performance in scenarios with a low number of images per class, thereby enhancing its potential. Furthermore, we maintain the improvement on cross-architectures and applications such as neural architecture search.

1 Introduction

Figure 1: The ATOM Framework utilizes inherent information to capture both context and location, resulting in significantly improved performance in dataset distillation. We display the performance of various components within the ATOM framework, showcasing a $5.8\%$ enhancement from the base distribution matching performance on CIFAR10 at IPC50. Complete numerical details can be found in Table 4.

Efficient deep learning has surged in recent years due to the increasing computational costs associated with training and inferencing pipelines [26, 63, 76, 53, 75, 51, 71, 52, 54, 2]. This growth can be attributed to the escalating complexity of model architectures and the ever-expanding scale of datasets. Despite the increasing computational burden, two distinct approaches have emerged as potential avenues for addressing this issue: the model-centric and data-centric approaches. The model-centric approach is primarily concerned with mitigating computational costs by refining the architecture of deep learning models. Techniques such as pruning, quantization, knowledge distillation, and architectural simplification are key strategies employed within this paradigm [26, 68, 71, 65, 29, 30, 50, 49]. In contrast, the data-centric approach adopts a different perspective, focusing on exploring and leveraging the inherent redundancy within datasets. Rather than modifying model architectures, this approach seeks to identify or construct a smaller dataset that retains the essential information necessary for maintaining performance levels. Coreset selection was a fairly adopted method for addressing this gap [47, 6, 4, 55, 60]. In particular works such as Herding [66] and K-Center [55] offered a heuristic-based approach to intelligently select an informative subset of data. However, as a heuristic-based method, the downstream performance is limited by the information contained solely in the subset. More recently, shapely data selection [17] found the optimal subset of data by measuring the downstream performance for every subset combination achievable in the dataset. However inefficient this may be, the downstream performance is still limited by the diversity of samples selected. therefore, Dataset Distillation (DD) [63] has emerged as a front-runner wherein a synthetic dataset can be learned.

Dataset distillation aims to distill large-scale datasets into a smaller representation, such that downstream models trained on this condensed dataset will retain competitive performance with those trained on the larger original one [63, 76, 7]. Recently, many techniques have been introduced to address this challenge, including gradient matching [76, 74, 38], feature/distribution matching [75, 51, 77], and trajectory matching [7, 14, 21]. However, many of these methods suffer from complex and computationally heavy distillation pipelines [76, 7, 21] or inferior performance [75, 51, 76]. A promising approach, DataDAM [51], effectively tackled the computational challenges present in prior distillation techniques by employing untrained neural networks, in contrast to bi-level optimization methods. However, despite its potential, DataDAM faced several significant limitations: (1) it obscured relevant class-content-based information existing channel-wise in intermediate layers; (2) it only achieved marginal enhancements on previous dataset distillation algorithms; and (3) it exhibited inferior cross-architecture generalization.

In this work, we introduce ATtentiOn Mixer, dubbed ATOM as an efficient dataset distillation pipeline that strikes an impressive balance between computational efficiency and superior performance. Drawing upon spatial attention matching techniques from prior studies like DataDAM [51], we expand our receptive field of information in the matching process. Our key contribution lies in mixing spatial information with channel-wise contextual information. Intuitively, different convolutional filters focus on different localizations of the input feature; thus, channel-wise attention aids in the distillation matching process by compressing and aggregating information from multiple regions as evident by the performance improvmenets displayed in Figure 1. ATOM not only combines localization and context, but it also produces distilled images that are more generalizable to various downstream architectures, implying that the distilled features are true representations of the original dataset. Moreover, our approach demonstrates consistent improvements across all settings on a comprehensive distillation test suite. In summary, the key contributions of this study can be outlined as follows:

[C1]: We provide further insight into the intricacies of attention matching, ultimately introducing the use of channel-wise attention matching for capturing a higher level of information in the feature-matching process. Our mixing module combines both spatial localization awareness of a particular class, with distinctive contextual information derived channel-wise.

[C2]: Empirically we show superior performance against previous dataset distillation methods including feature matching and attention matching works, without bi-level optimization on common computer vision datasets.

[C3]: We extend our findings by demonstrating superior performance in cross-architecture and neural architecture search. In particular, we provide a channel-only setting that maintains the majority of the performance while incurring a lower computational cost.

2 Related Works

Coreset Selection. Coreset selection, an early data-centric approach, aimed to efficiently choose a representative subset from a full dataset to enhance downstream training performance and efficiency. Various methods have been proposed in the past, including geometry-based approaches [1, 10, 55, 57, 66], loss-based techniques as mentioned in [59, 46], decision-boundary-focused methods [42, 16], bilevel optimization strategies [32, 33], and gradient-matching algorithms outlined in [43, 31]. Notable among them are Random, which randomly selects samples as the coreset; Herding, which picks samples closest to the cluster center; K-Center, which selects multiple center points to minimize the maximum distance between data points and their nearest center; and Forgetting, which identifies informative training samples based on learning difficulties [6, 4, 55, 59]. While these selection-based methods have shown moderate success in efficient training, they inherently possess limitations in capturing rich information. Since each image in the selected subset is treated independently, they lack the rich features that could have been captured if the diversity within classes had been considered. These limitations have motivated the emergence of dataset distillation within the field.

Figure 2: (a) An overview of the proposed ATOM framework. By mixing attention, ATOM is able to capture both spatial localization and class context. (b) Demonstration of the internal architecture for spatial- and channel-wise attention in the ATOM Module. The spatial-wise attention computes attention at specific locales through different filters, resulting in a matrix output, whereas the channel-wise attention calculates attention between each filter, naturally producing a vectorized output.

Dataset Distillation. Dataset distillation has emerged as a learnable method of synthesizing a smaller, information-rich dataset from a large-scale real dataset. This approach offers a more efficient training paradigm, commonly applied in various downstream applications such as continual learning [9, 51, 76, 20, 70], neural architecture search [27, 58], and federated learning [28, 69, 39, 40]. The seminal work, initially proposed by Wang et al. [63], introduced bilevel optimization, comprising an outer loop for learning the pixel-level synthetic dataset and an inner loop for training the matching network. Following this, several studies adopted surrogate objectives to tackle unrolled optimization problems in meta-learning. For example, gradient matching methods [76, 74, 38, 34, 15] learn images by aligning network gradients derived from real and synthetic datasets. Trajectory matching [7, 14, 11, 21] improves performance by minimizing differences in model training trajectories between original and synthetic samples. Meanwhile, feature matching strategies [75, 61, 51, 77, 73, 51] aim to align feature distributions between real and synthetic data within diverse latent spaces. Despite significant advancements in this field, methods still struggle to find a trade-off between the computational costs associated with the distillation pipeline and the model’s performance. A recent work, DataDAM [51], used spatial attention to improve the performance of feature-matching-based methods by selectively matching features based on their spatial attention scores. However, although this method operates without bilevel optimization, it only marginally improves performance on larger test suites. In this study, we delve deeper into the potential of attention-based methods and demonstrate superior performance compared to DataDAM and previous benchmarks across various computer vision datasets. Additionally, we achieve a lower computational cost compared to conventional attention-matching approaches by leveraging information in a channel-wise manner.

Attention Mechanism. Attention mechanisms have been widely adopted in deep learning to enhance performance across various tasks [3, 64, 72]. Initially applied in natural language processing [3], it has extended to computer vision, with global attention models [64] improving image classification and convolutional block attention modules [67] enhancing feature map selection. Additionally, attention aids model compression in knowledge distillation [72]. They are lauded for their ability to efficiently incorporate global contextual information into feature representations. When applied to feature maps, attention can take the form of either spatial or channel-based methods. Spatial methods focus on identifying the informative regions (”where”), while channel-based methods complementarily emphasize the informative features (”what”). Both spatial localization and channel information are crucial for identifying class characteristics. Recently, Sajedi et al. proposed DataDAM [51] to concentrate only on spatial attention, capturing class correlations within image localities for efficient training purposes. However, inspired by the inherent obfuscation of the content in the attention maps, we propose an Attention Mixer module that uses a unique combination of spatial and channel-wise attention to capture localization and information content.

3 Methodology

Given the larger source dataset $\mathcal{T}=\{(\bm{x}_{i},y_{i})\}_{i=1}^{|\mathcal{T}|}$ containing $|\mathcal{T}|$ real image-label pairs, we generate a smaller learnable synthetic dataset $\mathcal{S}=\{(\bm{s}_{j},y_{j})\}_{j=1}^{|\mathcal{S}|}$ with $|\mathcal{S}|$ synthetic image and label pairs. Following previous works [76, 74, 61, 51, 7], we use random sampling to initialize our synthetic dataset. For every class $k$, we obtain a batch of real and synthetic data ($B^{\mathcal{T}}_{k}$ and $B^{\mathcal{S}}_{k}$, respectively) and use a neural network $\phi_{\bm{\theta}}(\cdot)$ with randomly initialized weights $\bm{\theta}$ [22] to extract intermediate and output features. We illustrate our method in Figure 2 where an $L$-layer neural network $\phi_{\bm{\theta}}(\cdot)$ is used to extract features from the real and synthetic sets. The collection of feature maps from the real and synthetic sets can be expressed as $\phi_{\bm{\theta}}({\mathcal{T}}_{k})=[\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},1},\cdots,\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},L}]$ and $\phi_{\bm{\theta}}(\mathcal{S}_{k})=[\bm{f}^{\mathcal{S}_{k}}_{\bm{\theta},1},\cdots,\bm{f}^{\mathcal{S}_{k}}_{\bm{\theta},L}]$, respectively. The feature $\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l}$ comprises a multi-dimensional array within $\mathbb{R}^{|B^{\mathcal{T}}_{k}|\times C_{l}\times W_{l}\times H_{l}}$, obtained from the real dataset at the $l^{\text{th}}$ layer, where $C_{l}$ denotes the number of channels and $H_{l}\times W_{l}$ represents the spatial dimensions. Correspondingly, a feature $\bm{f}^{\mathcal{S}_{k}}_{\bm{\theta},l}$ is derived for the synthetic dataset.

We now introduce the Attention Mixer Module (ATOM) which generates attention maps for the intermediate features derived from both the real and synthetic datasets. Leveraging a feature-based map** function $A(\cdot)$, ATOM takes the intermediate feature maps as input and produces a corresponding attention map for each feature. Formally, we express this as: $A\big{(}\phi_{\bm{\theta}}({\mathcal{T}}_{k})\big{)}=[\bm{a}^{\mathcal{T}_{k}}_{\bm{\theta},1},\cdots,\bm{a}^{\mathcal{T}_{k}}_{\bm{\theta},L-1}]$ and $A(\phi_{\bm{\theta}}({\mathcal{S}}_{k}))=[\bm{a}^{\mathcal{S}_{k}}_{\bm{\theta},1},\cdots,\bm{a}^{\mathcal{S}_{k}}_{\bm{\theta},L-1}]$ for the real and synthetic sets, respectively. Previous works [51, 72] have shown that spatial attention, which aggregates the absolute values of feature maps across the channel dimension, can emphasize common spatial locations associated with high neuron activation. The implication of this is retaining the most informative regions, thus generating an efficient feature descriptor. In this work, we also consider the effect of channel-wise attention, which emphasizes the most significant information captured by each channel based on the magnitude of its activation. Since different filters explore different regions or locations of the input feature, channel-wise activation yields the best aggregation of the global information. Ultimately, we convert the feature map $\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l}$ of the $l^{\text{th}}$ layer into an attention map $\bm{a}^{\mathcal{T}_{k}}_{\bm{\theta},l}$ representing spatial or channel-wise attention using the corresponding map** functions $A_{s}(\cdot)$ or $A_{c}(\cdot)$ respectively. Formally, we can denote the spatial and channel-wise attention maps as:

$$A_{s}(\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l})=\sum_{i=1}^{C_{l}}\big{|}{(\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l})}_{i}\big{|}^{p_{s}},$$ (1)

$$A_{c}(\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l})=\sum_{i=1}^{H_{l}*W_{l}}\big{|}(\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l})^{\star}_{i}\big{|}^{p_{c}},$$ (2)

where, ${(\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l})}_{i}=\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l}(:,i,:,:)$ is the feature map of channel $i$ from the $l^{\text{th}}$ layer, and the power and absolute value operations are applied element-wise; meanwhile, the symbol ^⋆ flattens the feature map along the spatial dimension $\left((\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l})^{*}\in\mathbb{R}^{|B^{\mathcal{T}}_{k}|\times C_{l}\times W_{l}*H_{l}}\right)$, such that $(\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l})^{\star}_{i}=(\bm{f}^{\mathcal{T}_{k}}_{\bm{\theta},l})^{\star}(:,:,i)$. By leveraging both types of attention, we can better encapsulate the relevant information in the intermediate features, as investigated in Figure 4. Further, the effect of power parameters for spatial and channel-wise attention, i.e. $p_{s}$ and $p_{c}$ is studied in the Table 4.

Given our generated spatial and channel attention maps for the intermediate features, we apply standard normalization such that we can formulate a matching loss between the synthetic and real datasets. We denote our generalized loss $\mathcal{L}_{\text{ATOM}}$ as:

$$\displaystyle\mathop{\mathbb{E}}_{\bm{\theta}\sim P_{\bm{\theta}}}\bigg{[}\sum_{k=1}^{K}\sum_{l=1}^{L-1}\Big{\lVert}\displaystyle{\mathbb{E}}_{\mathcal{T}_{k}}\Big{[}\frac{\bm{z}^{\mathcal{T}_{k}}_{\bm{\theta},l}}{{\lVert\bm{z}^{\mathcal{T}_{k}}_{\bm{\theta},l}\rVert}_{2}}\Big{]}-\displaystyle\mathbb{E}_{{\mathcal{S}}_{k}}\Big{[}\frac{\bm{z}^{\mathcal{S}_{k}}_{\bm{\theta},l}}{{\lVert\bm{z}^{\mathcal{S}_{k}}_{\bm{\theta},l}\rVert}_{2}}\Big{]}\Big{\rVert}^{2}\bigg{]},$$ (3)

where, in the case of spatial attention, we denote $\bm{z}^{\mathcal{T}_{k}}_{\bm{\theta},l}=vec(\bm{a}^{\mathcal{T}_{k}}_{\bm{\theta},l})\in\mathbb{R}^{|B^{\mathcal{T}}_{k}|\times(W_{l}\times H_{l})}$ and $\bm{z}^{\mathcal{S}_{k}}_{\bm{\theta},l}=vec(\bm{a}^{\mathcal{S}_{k}}_{\bm{\theta},l})\in\mathbb{R}^{|B^{\mathcal{S}}_{k}|\times(W_{l}\times H_{l})}$ to represent the vectorized spatial attention map pairs at the $l^{\text{th}}$ layer for the real and synthetic datasets, respectively. Meanwhile, for channel-based attention, we have $\bm{z}^{\mathcal{T}_{k}}_{\bm{\theta},l}=vec(\bm{a}^{\mathcal{T}_{k}}_{\bm{\theta},l})\in\mathbb{R}^{|B^{\mathcal{T}}_{k}|\times(C_{l})}$ and $\bm{z}^{\mathcal{S}_{k}}_{\bm{\theta},l}=vec(\bm{a}^{\mathcal{S}_{k}}_{\bm{\theta},l})\in\mathbb{R}^{|B^{\mathcal{S}}_{k}|\times(C_{l})}$ to represent the flattened channel attention map pairs at the $l^{\text{th}}$ layer for the real and synthetic datasets, respectively. The parameter $K$ is the number of categories in a dataset, and $P_{\bm{\theta}}$ denotes the distribution of network parameters. We estimate the expectation terms in Equation 3 empirically if ground-truth data distributions are not available.

Following previous works [51, 75, 61, 77, 73], we leverage the features in the final layer to regularize our matching process. In particular, the features of the penultimate layer represent a high-level abstraction of information from the input images in an embedded representation and can thus be used to inject semantic information in the matching process [51, 48, 75, 19]. Thus, we employ $\mathcal{L}_{\text{MMD}}$ as described in [51, 75] out-of-the-box.

Finally, we learn the synthetic dataset by minimizing the following optimization problem using SGD optimizer:

$$\mathcal{S}^{*}=\operatorname*{arg\,min}_{\mathcal{S}}\>\big{(}\mathcal{L}_{\text{ATOM}}+\lambda\mathcal{L}_{\text{MMD}}\big{)},$$ (4)

where $\lambda$ is the task balance parameter inherited from [51]. In particular, we highlight that $\mathcal{L}_{\text{MMD}}$ brings semantic information from the final layer, while $\mathcal{L}_{\text{ATOM}}$ mixes the spatial and channel-wise attention information from the intermediate layers. Note that our approach assigns a fixed label to each synthetic sample and keeps it constant during training. A summary of the learning algorithm can be found in Algorithm 1.

Algorithm 1 Attention Mixer for Dataset Distillation

Input: Real training dataset $\mathcal{T}=\{(\bm{x}_{i},y_{i})\}_{i=1}^{|\mathcal{T}|}$
Required: Initialized synthetic samples for $K$ classes, Deep neural network $\phi_{\bm{\theta}}$ parameterized with $\bm{\bm{\theta}}$, Probability distribution over randomly initialized weights $P_{\bm{\theta}}$, Learning rate $\eta_{\mathcal{S}}$, Task balance parameter $\lambda$, Number of training iterations $I$.

1:Initialize synthetic dataset $\mathcal{S}$

2:for $i=1,2,\cdots,I$ do

3:     Sample $\bm{\theta}$ from $P_{\bm{\theta}}$

4:     Sample mini-batch pairs $B_{k}^{\mathcal{T}}$ and $B_{k}^{\mathcal{S}}$ from the real

5:       and synthetic sets for each class $k$

6:     Compute $\mathcal{L}_{\text{ATOM}}$ and $\mathcal{L}_{\text{MMD}}$

7:     Calculate $\mathcal{L}=\mathcal{L}_{\text{ATOM}}+\lambda\mathcal{L}_{\text{MMD}}$

8:     Update the synthetic dataset using $\mathcal{S}\leftarrow\mathcal{S}-\eta_{\mathcal{S}}\nabla_{\mathcal{S}}\mathcal{L}$

9:end for

Output: Synthetic dataset $\mathcal{S}=\{(\bm{s}_{i},y_{i})\}_{i=1}^{|\mathcal{S}|}$

4 Experiments

4.1 Experimental Setup

Datasets. Our method is evaluated on the CIFAR-10 and CIFAR-100 datasets [35], which maintain a resolution of 32 $\times$ 32, aligning with state-of-the-art benchmarks. Furthermore, we resize the Tiny ImageNet [37] datasets to 64 $\times$ 64 for additional experimentation. The supplementary materials provide more detailed dataset information.

Dataset	CIFAR-10			CIFAR-100			Tiny ImageNet
IPC	1	10	50	1	10	50	1	10	50
Ratio %	0.02	0.2	1	0.2	2	10	0.2	2	10
Random	$14.4_{\pm 2.0}$	$26.0_{\pm 1.2}$	$43.4\textpm_{\pm 1.0}$	$4.2\textpm_{\pm 0.3}$	$14.6_{\pm 0.5}$	$30.0_{\pm 0.4}$	$1.4_{\pm 0.1}$	$5.0_{\pm 0.2}$	$15.0_{\pm 0.4}$
Herding [66]	$21.5_{\pm 1.2}$	$31.6_{\pm 0.7}$	$40.4_{\pm 0.6}$	$8.3_{\pm 0.3}$	$17.3_{\pm 0.3}$	$33.7_{\pm 0.5}$	$2.8_{\pm 0.2}$	$6.3_{\pm 0.2}$	$16.7_{\pm 0.3}$
K-Center [55]	$21.5_{\pm 1.3}$	$14.7_{\pm 0.9}$	$27.0_{\pm 1.4}$	$8.4_{\pm 0.3}$	$17.3_{\pm 0.3}$	$30.5_{\pm 0.3}$	-	-	-
Forgetting [59]	$13.5_{\pm 1.2}$	$23.3_{\pm 1.0}$	$23.3_{\pm 1.1}$	$4.5_{\pm 0.2}$	$15.1_{\pm 0.3}$	-	$1.6_{\pm 0.1}$	$5.1_{\pm 0.2}$	$15.0_{\pm 0.3}$
DD†[63]	-	$36.8_{\pm 1.2}$	-	-	-	-	-	-	-
LD†[5]	$25.7_{\pm 0.7}$	$38.3_{\pm 0.4}$	$42.5_{\pm 0.4}$	$11.5_{\pm 0.4}$	-	-	-	-
DC [76]	$28.3_{\pm 0.5}$	$44.9_{\pm 0.5}$	$53.9_{\pm 0.5}$	$12.8_{\pm 0.3}$	$25.2_{\pm 0.3}$	$30.6_{\pm 0.6}$	$5.3_{\pm 0.1}$	$12.9_{\pm 0.1}$	$12.7_{\pm 0.4}$
DCC [38]	$32.9_{\pm 0.8}$	$49.4_{\pm 0.5}$	$61.6_{\pm 0.4}$	$13.3_{\pm 0.3}$	$30.6_{\pm 0.4}$	-	-	-	-
DSA [74]	$28.8_{\pm 0.7}$	$52.1_{\pm 0.5}$	$60.6_{\pm 0.5}$	$13.9_{\pm 0.3}$	$32.3_{\pm 0.3}$	$42.8_{\pm 0.4}$	$5.7_{\pm 0.1}$	$16.3_{\pm 0.2}$	$15.1_{\pm 0.2}$
DM [75]	$26.0_{\pm 0.8}$	$48.9_{\pm 0.6}$	$63.0_{\pm 0.4}$	$11.4_{\pm 0.3}$	$29.7_{\pm 0.3}$	$43.6_{\pm 0.4}$	$3.9_{\pm 0.2}$	$12.9_{\pm 0.4}$	$25.3_{\pm 0.2}$
GLaD [8]	$28.0_{\pm 0.8}$	$46.7_{\pm 0.5}$	$59.9_{\pm 0.7}$	-	-	-	-	-	-
CAFE [61]	$30.3_{\pm 1.1}$	$46.3_{\pm 0.6}$	$55.5_{\pm 0.6}$	$12.9_{\pm 0.3}$	$27.8_{\pm 0.3}$	$37.9_{\pm 0.3}$	-	-	-
CAFE+DSA [61]	$31.6_{\pm 0.8}$	$50.9_{\pm 0.5}$	$62.3_{\pm 0.4}$	$14.0_{\pm 0.3}$	$31.5_{\pm 0.2}$	$42.9_{\pm 0.2}$	-	-	-
VIG [41]	$26.5_{\pm 1.2}$	$54.6_{\pm 0.1}$	$35.6_{\pm 0.6}$	$17.8_{\pm 0.1}$	$29.3_{\pm 0.1}$	-	-	-	-
KIP [44]	$29.8_{\pm 1.0}$	$46.1_{\pm 0.7}$	$53.2_{\pm 0.7}$	$12.0_{\pm 0.2}$	$29.0_{\pm 0.3}$	-	-	-	-
MTT [7]	$31.9_{\pm 1.2}$	$56.4_{\pm 0.7}$	$65.9_{\pm 0.6}$	$13.8_{\pm 0.6}$	$33.1_{\pm 0.4}$	$42.9_{\pm 0.3}$	$6.2_{\pm 0.4}$	$17.3_{\pm 0.2}$	$26.5_{\pm 0.3}$
DAM [51]	$32.0_{\pm 1.2}$	$54.2_{\pm 0.8}$	$67.0_{\pm 0.4}$	$14.5_{\pm 0.5}$	$34.8_{\pm 0.5}$	$49.4_{\pm 0.3}$	$8.3_{\pm 0.4}$	$18.7_{\pm 0.3}$	$28.7_{\pm 0.3}$
ATOM (Ours)	$\mathbf{34.8}_{\pm 1.0}$	$\mathbf{57.9}_{\pm 0.7}$	$\mathbf{68.8}_{\pm 0.5}$	$\mathbf{18.1}_{\pm 0.4}$	$\mathbf{35.7}_{\pm 0.4}$	$\mathbf{50.2}_{\pm 0.3}$	$\mathbf{9.1}_{\pm 0.2}$	$\mathbf{19.5}_{\pm 0.4}$	$\mathbf{29.1}_{\pm 0.3}$
Full Dataset	$84.8_{\pm 0.1}$			$56.2_{\pm 0.3}$			$37.6_{\pm 0.4}$

Table 1: Comparison with previous dataset distillation methods on CIFAR-10, CIFAR-100 and Tiny ImageNet. The works DD^† and LD^† use AlexNet [36] for CIFAR-10 dataset. All other methods use ConvNet for training and evaluation. Bold entries are the best results.

Network Architectures. We employ a ConvNet architecture [18] for distillation, following prior studies. The default ConvNet comprises three convolutional blocks, each consisting of a 128-kernel 3 $\times$ 3 convolutional layer, instance normalization, ReLU activation, and 3 $\times$ 3 average pooling with a stride of 2. To accommodate the increased resolutions in Tiny ImageNet, we append a fourth convolutional block. Network parameters are initialized using normal initialization [22] in all experiments.

Evaluation Protocol. We evaluate the methods using standard measures from previous studies [75, 76, 61, 74, 51]. Five sets of synthetic images are generated from a real training dataset with 1, 10, and 50 images per class. Then, 20 neural network models are trained on each synthetic set using an SGD optimizer with a fixed learning rate of 0.01. Each experiment reports the mean and standard deviation values for 100 models to assess the efficacy of distilled datasets. Furthermore, computational costs are assessed by calculating run-time per step over 100 iterations, as well as peak GPU memory usage during 100 iterations of training.

Implementation Details. We use the SGD optimizer with a fixed learning rate of 1 to learn synthetic datasets containing 1, 10, and 50 IPCs over 8000 iterations with task balances ($\lambda$) set at 0.01. Previous works have shown that $p_{s}=4$ is sufficient for spatial attention matching [51]. As such we set our default case as: $p_{c}=p_{s}=4$. This is further ablated in Table 4. We adopt differentiable augmentation for both training and evaluating the synthetic set, following [76, 51]. For dataset reprocessing, we utilized the Kornia implementation of Zero Component Analysis (ZCA) with default parameters, following previous works [44, 7, 51]. All experiments are performed on a single A100 GPU with 80 GB of memory. Further hyperparameter details can be found in the supplementary materials.

Competitive Methods. In this paper, we compare the empirical results of ATOM on three computer vision datasets: CIFAR10/100 and TinyImageNet. We evaluate ATOM against four corset selection approaches and thirteen distillation methods for training set synthesis. The corset selection methods include Random selection [47], Herding [6, 4], K-Center [55], and Forgetting [60]. We also compare our approach with state-of-the-art distillation methods, including Dataset Distillation [63] (DD), Flexible Dataset Distillation [5] (LD), Dataset Condensation [76] (DC), Dataset Condensation with Contrastive (DCC) [38], Dataset Condensation with Differentiable Siamese Augmentation [74] (DSA), Distribution Matching [75] (DM), Deep Generative Priors (GLaD), Aligning Features [61] (CAFE), VIG [41], Kernel Inducing Points [44, 45] (KIP), Matching Training Trajectories [7] (MTT), and Attention Matching [51] (DAM).

4.2 Comparison with State-of-the-art Methods

Performance Comparison. In this section, we present a comparative analysis of our method against coreset and dataset distillation approaches. ATOM consistently outperforms these studies, especially at smaller distillation ratios, as shown in Table 1. Since the goal of dataset distillation is to generate a more compact synthetic set, we emphasize our significant performance improvements at low IPCs. We achieve almost $4\%$ improvement over the previous attention matching framework [51], DataDAM when evaluated on CIFAR-100 at IPC1. Notably, our performance on CIFAR-100 at IPC50 is 50.2% – that is nearly 90% of the baseline accuracy at a mere 10% of the original dataset. These examples motivate the development of dataset distillation works as downstream models can achieve relatively competitive performance with their baselines at a fraction of the training costs. Our primary objective in this study is to investigate the impact of channel-wise attention within the feature-matching process. Compared to prior attention-based and feature-based methodologies, our findings underscore the significance of channel-wise attention and the ATOM module, as validated also in the ablation studies in Figure 4.

Cross-architecture Generalization. In this section, we assess the generalization capacity of our refined dataset by training various unseen deep neural networks on it and then evaluating their performance on downstream classification tasks. Following established benchmarks [76, 75, 61, 51], we examine classic CNN architectures such as AlexNet [36], VGG-11 [56], ResNet-18 [23], and additionally, a standard Vision Transformer (ViT) [13]. Specifically, we utilize synthetic images learned from CIFAR-10 with IPC50 using ConvNet as the reference model and subsequently train the aforementioned networks on the refined dataset to assess their performance on downstream tasks. The results, as depicted in Table 2, indicate that ATOM demonstrates superior generalization across a spectrum of architectures. Notably, it achieves a significant performance boost of over 4% compared to the prior state-of-the-art on ResNet-18 [23]. This implies that the channel-wise attention mechanism effectively identifies features not only relevant to ConvNet but also to a wider range of deep neural networks, thereby enhancing the refined dataset with this discerned information.

	ConvNet	AlexNet	VGG-11	ResNet-18	ViT	Avg.
DC [76]	${53.9_{\pm 0.5}}$	${28.8_{\pm 0.7}}$	${38.8_{\pm 1.1}}$	${20.9_{\pm 1.0}}$	${30.1_{\pm 0.5}}$	${34.5_{\pm 0.8}}$
CAFE [61]	${62.3_{\pm 0.4}}$	${43.2_{\pm 0.4}}$	${48.8_{\pm 0.5}}$	${43.3_{\pm 0.7}}$	${22.7_{\pm 0.7}}$	${44.1_{\pm 0.5}}$
DSA [74]	$60.6_{\pm 0.5}$	$53.7_{\pm 0.6}$	$51.4_{\pm 1.0}$	$47.8_{\pm 0.9}$	${43.3_{\pm 0.4}}$	${51.4_{\pm 0.7}}$
DM [75]	$63.0_{\pm 0.4}$	$60.1_{\pm 0.5}$	$57.4_{\pm 0.8}$	$52.9_{\pm 0.4}$	${45.2_{\pm 0.4}}$	${55.7_{\pm 0.5}}$
KIP [44]	$56.9_{\pm 0.4}$	$53.2_{\pm 1.6}$	$53.2_{\pm 0.5}$	$47.6_{\pm 0.8}$	${18.3_{\pm 0.6}}$	${45.8_{\pm 0.8}}$
MTT [7]	$66.2_{\pm 0.6}$	$43.9_{\pm 0.9}$	$48.7_{\pm 1.3}$	$60.0_{\pm 0.7}$	${47.7_{\pm 0.6}}$	${53.3_{\pm 0.8}}$
DAM [51]	$67.0_{\pm 0.4}$	$63.9_{\pm 0.9}$	$64.8_{\pm 0.5}$	$60.2_{\pm 0.7}$	${48.2_{\pm 0.8}}$	${60.8_{\pm 0.7}}$
ATOM (Ours)	$\mathbf{68.8}_{\pm 0.4}$	$\mathbf{64.1}_{\pm 0.7}$	$\mathbf{66.4}_{\pm 0.6}$	$\mathbf{64.5}_{\pm 0.6}$	$\mathbf{49.5}_{\pm 0.7}$	$\mathbf{62.7}_{\pm 0.6}$

Table 2: Cross-architecture testing performance (%) on CIFAR-10 with 50 images per class. The ConvNet architecture is employed for distillation. Bold entries are the best results.

Method	Run Time (Sec.)			GPU memory (MB)
	IPC1	IPC10	IPC50	IPC1	IPC10	IPC50
DC [76]	$0.16_{\pm 0.01}$	$3.31_{\pm 0.02}$	$15.74_{\pm 0.10}$	3515	3621	4527
DSA [74]	$0.22_{\pm 0.02}$	$4.47_{\pm 0.12}$	$20.13_{\pm 0.58}$	3513	3639	4539
DM [75]	$0.08_{\pm 0.02}$	$0.08_{\pm 0.02}$	$0.08_{\pm 0.02}$	3323	3455	3605
MTT [7]	$0.36_{\pm 0.23}$	$0.40_{\pm 0.20}$	OOM	2711	8049	OOM
DAM [51]	$0.09_{\pm 0.01}$	$0.08_{\pm 0.01}$	$0.16_{\pm 0.04}$	3452	3561	3724
$\textbf{ATOM}^{\dagger}$ (Ours)	$0.08_{\pm 0.02}$	$0.08_{\pm 0.02}$	$0.13_{\pm 0.03}$	3152	3263	4151
ATOM (Ours)	$0.10_{\pm 0.02}$	$0.10_{\pm 0.01}$	$0.17_{\pm 0.02}$	3601	4314	5134

Table 3: Comparisons of training time and GPU memory usage for prior dataset distillation methods. Run time is averaged per step over 100 iterations, while GPU memory usage is reported as peak memory during the same 100 iterations of training on an A100 GPU for CIFAR-10. Methods that surpass the GPU memory threshold and fail to run are denoted as OOM (out-of-memory). $\textbf{ATOM}^{\dagger}$ represents our method with on-channel attention, hence offering a better tradeoff in computational complexity.

Distillation Cost Analysis. In this section, we delve into an examination of the training costs required for the distillation process. Although the main goal of dataset distillation is to reduce training costs across different applications such as neural architecture search and continual learning, the distillation technique itself must be efficient, enabling smooth operation on consumer-grade hardware. Approaches such as DC, DSA, and MTT introduce additional computational overhead due to bi-level optimization and training an expert model. In contrast, our method, akin to DM and DAM, capitalizes on randomly initialized networks, obviating the need for training and thereby reducing the computational cost per step involved in the matching stage. As illustrated in Table 3 utilizing solely the channel-based $\texttt{ATOM}^{\dagger}$ decreases the computational burden of matching compared to the default ATOM configuration. This efficiency is crucial, as channel-wise attention offers a more effective distillation process while maintaining superior performance (refer to Figure 4).

Convergence Speed Analysis. In  Figure 3, we plot the downstream testing accuracy evolution for the synthetic images on CIFAR10 IPC50. Comparing with previous methods, DM [75] and DataDAM [51], we can explicitly see an improvement in convergence speed and a significantly higher steady state achieved with the ATOM framework. Our included convergence analysis supports the practicality of our method and the consistency to which we outperform previous baselines.

Figure 3: Test accuracy evolution of synthetic image learning on CIFAR10 with IPC50 for ATOM (ours), DM [75] and DataDAM [51].

4.3 Ablation Studies and Analysis

Figure 4: Sample learned synthetic images for CIFAR-10/100 (32$\times$32 resolution) IPC10 and TinyImageNet (64$\times$64 resolution) IPC 1.

Evaluation of loss components in ATOM. In Table 4, we evaluate the effect of different attention-matching mechanisms with respect to pure feature matching in intermediate layers and distribution matching in the final layer ($\mathcal{L}_{\text{MMD}}$). The results clearly demonstrate that attention-matching improves the performance of the distillation process. In particular, the attention-matching process improves feature matching by $8.0\%$. Further, it seems that channel attention is able to capture the majority of relevant information from the intermediate features, as evidenced by an improvement of over $1.5\%$ from spatial attention matching. Ultimately, this provides an incentive to favor channel attention in the distillation process.

$\mathcal{L}_{\text{MMD}}$	Feature Map	Spatial Atn.	Channel Atn.	Performance (%)
✓	-	-	-	$63.0_{\pm 0.4}$
	✓	-	-	$60.8_{\pm 0.6}$
✓	-	✓	-	$67.0_{\pm 0.7}$
✓	-	-	✓	$68.6_{\pm 0.3}$
✓	-	✓	✓	$\mathbf{68.8}_{\pm 0.5}$

Table 4: Evaluation of loss components and attention components in ATOM using CIFAR-10 with IPC50.

Evaluating attention balance in ATOM. In this section, we evaluate the balance between spatial and channel-wise attention through the power value $p$. Referencing  Equation 1 and  Equation 2, modulating the values of $p_{s}$ and $p_{c}$ ultimately affects the balance of spatial and channel-wise attention in $\mathcal{L}_{\text{ATOM}}$. In Table 5, we examine the impact of different exponentiation powers $p$ in the attention-matching mechanisms. Specifically, we conduct a grid-based search to investigate how varying the exponentiation of spatial ($p_{s}$) and channel ($p_{c}$) attention influences subsequent performance. Our findings reveal that optimal performance (nearly $1\%$ improvement over our default) occurs when the exponentiation for channel attention significantly exceeds that of spatial attention. This suggests that assigning a higher exponential value places greater emphasis on channel-attention matching over spatial-wise matching. This aligns with our observations from the loss component ablation, where channel-wise matching was found to encapsulate the majority of information within the feature map. Consequently, we deduce that prioritizing channel-wise matching will enhance downstream performance outcomes.

Channel Attention $p_{c}$	Spatial Attention $p_{s}$
	1	2	4	8
1	57.4%	57.5%	57.0%	56.2%
2	58.2%	57.5%	57.2%	56.3%
4	58.4%	58.5%	57.9%	57.6%
8	58.8%	58.7%	58.2%	57.8%

Table 5: Evaluation of power values in the spatial and channel attention computations for $\mathcal{L}_{\text{ATOM}}$ using CIFAR-10 with IPC10.

Visualization of Synthetic Images. We include samples of our distilled images in Figure 4. The images appear to be interleaved with artifacts that assimilate the background and object information into a mixed collage-like appearance. The synthetic images effectively capture the correlation between background and object elements, suggesting their potential for generalizability across various architectures, as empirically verified in Table 2. Additional visualizations are available in the supplementary material.

4.4 Applications

Neural Architecture Search. In Table Table 6 we leverage our distilled synthetic datasets as proxy sets to accelerate Neural Architecture Search. In line with previous state-of-the-art, [51, 76, 74], we outline our architectural search space, comprising 720 ConvNets on the CIFAR-10 dataset. We commence with a foundational ConvNet and devise a consistent grid, varying in depth $D\in$ {1, 2, 3, 4}, width $W\in$ {32, 64, 128, 256}, activation function $A\in$ {Sigmoid, ReLU, LeakyReLU}, normalization technique $N\in$ {None, BatchNorm, LayerNorm, InstanceNorm, GroupNorm}, and pooling operation $P\in$ {None, MaxPooling, AvgPooling}. Additionally, we benchmark our approach against several state-of-the-art methods, including Random, DSA [76], DM [75], CAFE [61], DAM [51], and Early-Stop**. Our method demonstrates superior performance, accompanied by a heightened Spearman’s correlation (0.75), thereby reinforcing the robustness of ATOM and its potential in neural architecture search.

	Random	DSA	DM	CAFE	DAM	ATOM	Early-stop**	Full Dataset
Performance (%)	88.9	87.2	87.2	83.6	89.0	88.9	88.9	89.2
Correlation	0.70	0.66	0.71	0.59	0.72	0.75	0.69	1.00
Time cost (min)	206.4	206.4	206.6	206.4	206.4	206.4	206.2	5168.9
Storage (imgs)	500	500	500	500	500	500	$5\times 10^{4}$	$5\times 10^{4}$

Table 6: Neural architecture search on CIFAR-10 with IPC50.

5 Limitations

Many studies in dataset distillation encounter a constraint known as re-distillation costs [62, 24, 25]. This limitation becomes apparent when adjusting the number of images per class (IPC) or the distillation ratios. Like most other distillation methods, our approach requires re-distillation on the updated setting configuration, which limits flexibility regarding configuration changes and storage allocation. Additionally, we observed in Table 2 that dataset distillation methods often struggle with generalizing to transformer architectures. Despite ATOM outperforming other methods, there is still a noticeable performance drop compared to convolutional neural networks. This suggests that the effectiveness of transformers for downstream training might be constrained by the distilled data.

6 Conclusion

In this work, we introduced an Attention Mixer (ATOM) for efficient dataset distillation. Previous approaches have struggled with marginal performance gains, obfuscating channel-wise information, and high computational overheads. ATOM addresses these issues by effectively combining information from different attention mechanisms, facilitating a more informative distillation process with untrained neural networks. Our approach utilizes a broader receptive field to capture spatial information while preserving distinct content information at the channel level, thus better aligning synthetic and real datasets. By capturing information across intermediate layers, ATOM facilitates multi-scale distillation. We demonstrated the superior performance of ATOM on standard distillation benchmarks and its favorable performance across multiple architectures. We conducted several ablative studies to justify the design choices behind ATOM. Furthermore, we applied our distilled data to Neural Architecture Search, showing a superior correlation with the real large-scale dataset. In the future, we aim to extend attention mixing to various downstream tasks, including image segmentation and localizations. We also hope to address limitations of ATOM, such as re-distillation costs and cross-architecture generalizations on transformers.

Supplementary Material

7 Implementation Details

7.1 Datasets

We conducted experiments on three main datasets: CIFAR10/100 [35] and TinyImageNet [37]. These datasets are considered single-label multi-class; hence, each image has exactly one class label. The CIFAR10/100 are conventional computer vision benchmarking datasets comprising 32$\times$32 colored natural images. They consist of 10 coarse-grained labels (CIFAR10) and 100 fine-grained labels (CIFAR100), each with 50,000 training samples and 10,000 test samples. The CIFAR10 classes include ”Airplane”, ”Car”, ”Bird”, ”Cat”, ”Deer”, ”Dog”, ”Frog”, ”Horse”, ”Ship”, and ”Truck”. The TinyImageNet dataset, a subset of ImageNet-1K [12] with 200 classes, contains 100,000 high-resolution training images and 10,000 test images resized to 64$\times$64. The experiments on these datasets make up the benchmarking for many previous dataset distillation works [76, 61, 51, 7, 8, 78].

7.2 Dataset Pre-processing

We applied the standardized preprocessing techniques to all datasets, following the guidelines provided in DM [75] and DataDAM [51]. Following previous works, we apply the default Differentiable Siamese Augmentation (DSA) [74] scheme during distillation and evaluation. Specifically for the CIFAR10/100 datasets, we integrated Kornia zero-phase component analysis (ZCA) whitening, following the parameters outlined in [7, 51]. Similar to DataDAM [51], we opted against ZCA for TinyImagenet due to the computational bottlenecks associated with full-scale ZCA transformation on a larger dataset with double the resolution. Note that we visualized the distilled images by directly applying the inverse transformation based on the corresponding data pre-processing, without any additional modifications.

7.3 Hyperparameters

Our method conveniently introduces only one additional hyperparameter: the power term in channel attention, i.e. $p_{c}$. All the other hyperparameters used in our method are directly inherited from the published work, DataDAM [51]. Therefore, we include an updated hyperparameter table in Table 7 aggregating our power term with the remaining pre-set hyperparameters. In the main paper, we discussed the effect of power terms on both channel- and spatial-wise attention and ultimately found that higher channel attention paired with lower spatial attention works best. However, our default, as stated in the main draft, is $p_{c}=p_{s}=4$. Regarding the distillation and train-val settings, we use the SGD optimizer with a learning rate of 1.0 for learning the synthetic images and a learning rate of 0.01 for training neural network models (for downstream evaluation). For CIFAR10/100 (low-resolution), we use a 3-layer ConvNet; meanwhile, for TinyImagenet (medium-resolution), we use a 4-layer ConvNet, following previous works in the field [75, 51, 7]. Our batch size for learning the synthetic images was set to 128 due to the computational overhead of a larger matching set.

7.4 Neural Architecture Search Details

Following previous works [51, 76, 74, 75], we define a search space consisting of 720 ConvNets on the CIFAR10 dataset. Models are evaluated on CIFAR10 using our IPC 50 distilled set as a proxy under the neural architecture search (NAS) framework. The architecture search space is constructed as a uniform grid that varies in depth $D\in$ {1, 2, 3, 4}, width $W\in$ {32, 64, 128, 256}, activation function $A\in$ {Sigmoid, ReLu, LeakyReLu}, normalization technique $N\in$ {None, BatchNorm, LayerNorm, InstanceNorm, GroupNorm}, and pooling operation $P\in$ {None, MaxPooling, AvgPooling} to create varying versions of the standard ConvNet. These candidate architectures are then evaluated based on their validation performance and ranked accordingly. In the main paper, Table 6 measures various costs and performance metrics associated with each distillation method. Overall distillation improves the computational cost; however, ATOM achieves the highest correlation, which is by far the most “important“ metric in this NAS search, as it indicates that our proxy set best estimates the original dataset.

8 Additional Visualizations.

We include additional visualizations of our synthetic datasets in Figure 5, Figure 6, Figure 7. The first two represent CIFAR10/100 at IPC 50, while the third depicts TinyImageNet at IPC 10. Our images highly exhibit learned artifacts from the distillation process that are, in turn, helpful during downstream classification tasks.

Figure 5: Distilled Image Visualization: CIFAR-10 dataset with IPC 50.

Figure 6: Distilled Image Visualization: CIFAR-100 dataset with IPC 50.

Figure 7: Distilled Image Visualization: TinyImageNet dataset with IPC 10.

Hyperparameters			Options/	Value
Category	Parameter Name	Description	Range
Optimization	Learning Rate $\bm{\eta_{\mathcal{S}}}$ (images)	Step size towards global/local minima	$(0,10.0]$	IPC $\leq$ 50: $1.0$
				IPC $>$ 50: $10.0$
	Learning Rate $\bm{\eta_{\mathcal{\bm{\theta}}}}$ (network)	Step size towards global/local minima	$(0,1.0]$	$0.01$
	Optimizer (images)	Updates synthetic set to approach global/local minima	SGD with	Momentum: $0.5$
			Momentum	Weight Decay: $0.0$
	Optimizer (network)	Updates model to approach global/local minima	SGD with	Momentum: $0.9$
			Momentum	Weight Decay: $5e-4$
	Scheduler (images)	-	-	-
	Scheduler (network)	Decays the learning rate over epochs	StepLR	Decay rate: $0.5$
				Step size: $15.0$
	Iteration Count	Number of iterations for learning synthetic data	$[1,\infty)$	8000
Loss Function	Task Balance $\lambda$	Regularization Multiplier	$[0,\infty)$	Low Resolution: $0.01$
				High Resolution: $0.02$
	Spatial Power Value $p_{s}$	Exponential power for amplification of spatial attention	$[1,\infty)$	4
	Channel Power Value $p_{c}$	Exponential power for amplification of channel attention	$[1,\infty)$	4
	Loss Configuration	Type of error function used to measure distribution discrepancy	-	Mean Squared Error
	Normalization Type	Type of normalization used in the SAM module on attention maps	-	L2
DSA Augmentations	Color	Randomly adjust (jitter) the color components of an image	brightness	1.0
			saturation	2.0
			contrast	0.5
	Crop	Crops an image with padding	ratio crop pad	0.125
	Cutout	Randomly covers input with a square	cutout ratio	0.5
	Flip	Flips an image with probability p in range:	$(0,1.0]$	$0.5$
	Scale	Shifts pixels either column-wise or row-wise	scaling ratio	$1.2$
	Rotate	Rotates image by certain angle	$0^{\circ}-360^{\circ}$	$[-15^{\circ},+15^{\circ}]$
Encoder Parameters	Conv Layer Weights	The weights of convolutional layers	$\mathbb{R}$ bounded by kernel size	Uniform Distribution
	Activation Function	The non-linear function at the end of each layer	-	ReLU
	Normalization Layer	Type of normalization layer used after convolutional blocks	-	InstanceNorm

Table 7: Hyperparameters Details – boilerplate obtained from DataDAM [51].