TLDR: This research introduces a novel method for self-supervised dataset distillation that significantly reduces the size of training datasets while maintaining high model performance and generalizability. It addresses limitations of previous methods by proposing innovative parameterization of images and representations using low-dimensional bases, and by incorporating predefined data augmentations with lightweight approximation networks to avoid gradient bias and reduce storage. Extensive experiments show that this approach outperforms existing baselines across various datasets and network architectures, demonstrating improved efficiency and transfer learning capabilities.
In the world of deep learning, the demand for ever-larger datasets like ImageNet and LAION has become a double-edged sword. While these massive datasets are crucial for training powerful models, they also bring significant challenges, primarily in the form of prohibitive training costs and computational expenses. This issue is particularly acute in scenarios requiring repeated training, such as continual learning or neural architecture search, which further inflates the financial burden.
To address this, two main strategies have emerged: coreset selection and dataset distillation (DD). Coreset selection identifies the most representative samples in a dataset, allowing models trained on these subsets to perform well. Dataset distillation, on the other hand, creates a synthetic, highly compact dataset optimized to achieve performance comparable to models trained on the full original dataset. While coreset selection is faster, dataset distillation generally yields better training outcomes despite longer processing times.
Historically, most dataset distillation efforts have focused on supervised datasets, meaning they rely heavily on labeled data. This approach often leads to models that emphasize class-specific features, potentially overfitting to particular models and limiting their ability to generalize across different architectures or transfer to new tasks. In contrast, self-supervised learning (SSL) has shown great promise in mapping images to generalized representations that transfer effectively to various downstream applications. This makes condensing these generalized representations a compelling avenue for creating distilled datasets with superior cross-architecture generalizability and improved transferability.
A recent pioneering work in self-supervised dataset distillation is KRR-ST (Kernel Ridge Regression on Self-Supervised Target). KRR-ST aims to distill image and representation pairs from unlabeled datasets, allowing a new model trained on this distilled set to mimic the self-supervised model trained on the entire unlabeled dataset. This new model can then serve as a good starting point for fine-tuning on other tasks. However, KRR-ST faces two significant limitations: it requires substantial storage for distilled images and their paired representations, and it bypasses random data augmentation during optimization due to incompatibility with the bilevel optimization process, despite augmentation being a critical component in self-supervised learning.
This new research introduces several novel techniques to overcome these challenges, significantly boosting self-supervised dataset distillation. The core of their approach lies in three key areas:
Innovative Parameterization for Compactness
The first major contribution is an innovative parameterization method for both images and their representations. Instead of storing each distilled image and representation independently, which can lead to redundancy, the authors propose representing them using distinct low-dimensional bases and coefficients. This idea is inspired by techniques like Eigenface, which uses principal components to efficiently represent faces. For images, the bases are initialized using the principal components of the original unlabeled dataset. Similarly, for representations, bases are initialized from the principal components of the representations extracted by the teacher model from the real data. This method not only compresses the data significantly but also improves performance by focusing on the most crucial characteristics.
Predetermined Augmentation and Approximation Networks
The second crucial innovation addresses the issue of data augmentation. While KRR-ST avoids random augmentations due to gradient bias in bilevel optimization, this new method reintroduces the benefits of augmentation by predetermining all augmentations used in the distillation process. This eliminates randomness and allows for the recording of corresponding augmented representations. To prevent this from consuming excessive memory, the authors introduce lightweight “approximation networks.” These networks learn to predict the shift in representation that occurs when an unaugmented image is transformed into its augmented views. This means only the approximation networks and the representations of the unaugmented distilled images need to be stored, leading to much more compact distilled sets.
Also Read:
- Advancing Deep Subspace Clustering with Mini-Batch Training and Memory Banks
- Teaching AI to Craft Smarter Digital Deceptions: A Knowledge Distillation Approach
Optimization and Evaluation
The proposed framework follows a bilevel optimization scheme, similar to KRR-ST, but with modifications to incorporate the new parameterization and augmentation techniques. An inner loop trains a feature extractor on the augmented distilled data, while an outer loop optimizes the bases and coefficients to ensure the feature extractor’s output aligns with that of a pre-trained teacher model on the real dataset. After this optimization, the approximation networks are trained to accurately predict representation shifts caused by augmentations.
The effectiveness of the distilled dataset is evaluated using a standard linear evaluation scheme common in self-supervised learning. A new feature extractor is trained on the distilled dataset, then frozen and coupled with a linear classifier. This classifier is trained on a supervised downstream task, and its performance indicates the quality of the distilled dataset and the generalized features learned. The experiments were conducted on various datasets, including CIFAR100, TinyImageNet, and ImageNet, and evaluated across different network architectures like CNN, VGG11, ResNet18, AlexNet, MobileNet, and Vision Transformer (ViT).
The results consistently demonstrate the superiority of this new approach. It significantly outperforms existing baselines, including KRR-ST and several supervised dataset distillation methods, in terms of distillation efficiency, cross-architecture generalization, and transfer learning performance. The ablation studies further confirm that both the parameterization and the predefined augmentation with approximation networks contribute substantially to these improvements. The research also highlights the critical role of proper initialization for the bases and coefficients, showing that using principal components from the real dataset leads to significantly better performance.
In conclusion, this work presents a powerful self-supervised dataset distillation method that effectively reduces dataset size while preserving model performance across diverse architectures. By leveraging innovative parameterization, predefined augmentations, and approximation networks, it creates highly compact and generalizable distilled datasets, addressing key limitations of prior work. You can read the full paper here.
