Meta-DM: Applications of Diffusion Models on Few-Shot Learning

The FSL problem is commonly defined as the $N$-way $K$-shot problem: given a base dataset $\mathcal{D}_{base}=\{(x_{i},y_{i})\}$, the goal is to train a model that can effectively solve the downstream few-shot task $\mathcal{T}$, which consists of a support set $\mathcal{S}=\{(x_{s},y_{s})\}^{N\times K}_{s=1}$ for adaptation and a query set $\mathcal{Q}=\{x_{q}\}^{Q}_{q=1}$ for prediction, where $y_{s}$ is the class label of image $x_{s}$.

The Prototypical Network is considered as one of the most classical solutions for FSL problems [1]. The main idea of Prototypical Networks is to find a feature prototype for each class by computing the mean vector of its embedded support points:

For a query point $x$, Prototypical Networks compute the Euclidean distance between $x$ and each feature prototype, and then produce a distribution over all classes by applying a softmax function to the distances:

The objective of training is to minimize the negative logarithm of the true class probability $k$: $J(\phi)=-\log p_{\phi}(y=k|x)$, thereby enhancing the classifier’s recognition capability.

P>M>F stands for the pipeline consisting of Pre-training, Meta-learning and Fine-tuning [11], each of which can be instantiated with different pre-training algorithm or backbone architecture. For example, Hu et al. uses DINO [23] for pre-training and Prototypical Networks with ViT [24] for meta-learning, resulting in fairly good performance on standard benchmarks. Additionally, P>M>F includes optional fine-tuning modules for cross-domain FSL tasks.

The definition of unsupervised FSL is similar to that of supervised FSL, with the difference being that only unlabeled data $\mathcal{D}_{base}=\{x_{i}\}$ is available for training unsupervised FSL models.

Meta-GMVAE (Meta-Gaussian Mixture Variational Autoencoders) [22] is a pioneering unsupervised FSL algorithm based on VAE (Variational Autoencoders) [25] and set-level variational inference using self-attention [26]. The core idea of Meta-GMVAE is to use multimodal prior distributions, a mixture of Gaussians, considering that each modality corresponds to a class in unsupervised FSL tasks. It uses the EM (Expectation-Maximization) algorithm to optimize the parameters of Gaussian Mixture Model.

UniSiam is an advanced unsupervised FSL algorithm that achieves state-of-the-art performance on several benchmarks [13]. Unlike in supervised FSL, where the goal is to maximize the mutual information $I(Z;Y)$ between the representation $Z$ and the label $Y$, UniSiam aims to maximize the mutual information $I(Z;X)$ between the representation $Z$ and the data $X$. This approach enables the model to learn a meaningful representation of each class, while also avoiding overfitting to the base classes.

In this process, noise from a standard Gaussian distribution is continuously added to the input data. Given a sample $s_{0}$, noise is iteratively added to it to obtain $s_{1},s_{2},\ldots,s_{T-1},s_{T}$. As the number of iterations tends to infinity, the resulting sequence of samples tends towards an isotropic Gaussian distribution.

In contrast to the forward process, the reverse process involves continuously filtering out noise and recovering data from the noisy data. As previously defined, we begin by sampling from $s_{T}$, and then gradually backtrack to $s_{T-1},s_{T-2},\cdots,s_{0}$.

The diffusion model is designed to learn the task of removing a small fraction of noise from $s_{t}$ to obtain $s_{t-1}$, where $t$ denotes any timestep during the process. To achieve this, the model employs a function $\epsilon_{\theta}(s_{t},t)$ [27] that predicts the noise component of $s_{t}$. The training samples consist of the data $s_{t}$ at a given moment, the timestep $t$, and the noise $\epsilon$, all of which can be sampled during the forward process. The training objective is to minimize the mean square error loss between the predicted noise $\epsilon_{\theta}(s_{t},t)$ and the actual noise $\epsilon$ [18], as shown in the equation below:

The miniImageNet dataset, originally proposed by Vinyals et al. [2], is a subset of the widely used ImageNet dataset [30]. It consists of 60,000 color images in either their original size or in a processed size of 84$\times$84 pixels. The dataset is divided into 100 classes, with 600 images in each class. For our experiments, we follow the data splits proposed by Ravi and Larochelle [7], where 64 classes are used for training, 16 classes for validation, and 20 classes for testing.

The tieredImageNet dataset, proposed by Ren et al. [8], is a larger subset of ImageNet [30] than miniImageNet. It comprises 608 classes, each with about 1,300 images. These classes are further grouped into 34 higher-level categories, which are divided into 20 categories for training, 6 categories for validation, and 8 categories for testing. Unlike other datasets, all classes in each category are used in only one stage, which ensures that all training classes are sufficiently fine-grained to differ from the test classes.

The Meta-Dataset, proposed by Triantafillou et al. [31], is a comprehensive benchmark dataset designed specifically for FSL. The dataset is constructed by combining images from 10 different datasets, including Omniglot, Aircraft, CUB-200-2011, etc. The Meta-Dataset is much larger than any previous FSL dataset, and each component image may have different types, styles, and sizes, which places higher demands on the generalizability of the FSL model [31].

To allow the classifier to learn more features, we employ the full Meta-DM module on Prototypical Networks, which have a relatively simple backbone and do not use external data for pre-training. Specifically, for each original image, we generate a "good" sample and a "bad" sample, and assign the "good" sample to the original class and the "bad" sample to the "extra" classes. To generate the "good" samples, we set the strength of the diffusion model to 0.05, while for the "bad" samples, we set it to 0.2. Our approach is compared against Prototypical Networks without Meta-DM and several other classical supervised FSL algorithms, with the results presented in Table 1. For more training details, please see our supplemental material.

As discussed in Section 2.1, there are multiple pre-training algorithms that can be used to instantiate P>M>F, such as DINO [23], BEiT [32], and CLIP [33]. We follow Hu et al. and use DINO for self-supervised pre-training. P>M>F also offers a variety of advanced backbone architectures, including ResNet-50 [34] and Vit [24], which have stronger feature extraction capabilities. Given that the classifier is expected to have learned sufficient intra-class features during pre-training, we use a simplified version of Meta-DM for P>M>F. Specifically, we randomly select 5 images out of the 600 in each class and generate only "bad" samples based on them. All "bad" samples are assigned to an "extra" class, and the strength for "bad" samples is still 0.2. Since we train and test on the same dataset, fine-tuning is not necessary. Table 2 presents the results of Meta-DM $+$ P>M>F with ResNet-50 and Vit-Small/Vit-Base as backbones. In the above two experiments, Meta-DM demonstrate its ability to achieve broad improvements in supervised FSL problems. More importantly, it is able to provide considerable enhancements to state-of-the-art algorithms.

To gain the performance of Meta-DM on cross-domain tasks, we train the model on miniImageNet and then test it on several subsets of the Meta-Dataset. We follow Hu et al. by first automatically selecting the learning rate during the model deployment phase and then fine-tuning the feature backbone through several gradient steps [11]. Meta-DM is only used in the training stage. We compare the results of cross-domain tasks with and without Meta-DM in Table 3. The results show that Meta-DM can also achieve a slight improvement on cross-domain tasks.

In this experiment, we utilize the complete Meta-DM module that is used on Prototypical Networks. To ensure accurate comparison, we only report the performance results on miniImageNet, which includes 5-way, 1-shot / 5-shot / 20-shot / 50-shot settings. The performance results are shown in Table 4.

In the original UniSiam by Lu et al., the 224 $\times$ 224 image input size and two effective data augmentation methods: RandomVerticalFlip and RandAugment [15], are used in order to improve the model’s ability to extract features from scratch [13]. Therefore, we utilize the simplified Meta-DM module as in P>M>F, with the only difference being that the strength of the diffusion model is adjusted to 0.3 as the image input size increases. For accurate comparison, we report the performance of Meta-DM $+$ UniSiam with ResNet-18 / ResNet-34 / ResNet-50 as backbones and miniImageNet / tieredImageNet as benchmarks. To further push the limits of unsupervised FSL problems, we also conduct standard knowledge distillation [40] on each of our models, with the teacher models being the previous models we trained with Meta-DM $+$ UniSiam and ResNet-50. The results are shown in Table 5. Our experiments demonstrate that Meta-DM also achieves state-of-the-art performance on unsupervised FSL. Specifically, it exhibits considerable improvements on miniImageNet and slight improvements on tieredImageNet. This suggests that Meta-DM is better suited for datasets with fewer classes and fewer samples per class, which facilitates the learning of new features by the classifier and ultimately leads to better performance. Furthermore, the increase in computational power requirement resulting from the generating a small number of "bad" samples is negligible compared to the original experiment.