MetaDiff: Meta-Learning with Conditional Diffusion for Few-Shot Learning

Few-shot learning (FSL) is a challenging task, which aims to recognize novel classes with few examples (Chen et al. 2021, 2019). To address this problem, meta-learning is proposed, which aims to learn to quickly learn novel tasks with few examples (Flennerhag et al. 2019; Zhu and Koniusz 2023). The core idea is learning task-agnostic meta knowledge from a large number of similar tasks and then leveraging it to assist the learning of novel tasks. From the type of meta-knowledge, these existing meta-learning methods can be roughly grouped into three groups. The metric-based methods (Snell et al. 2017; Vinyals et al. 2016; Zhang et al. 2021a, 2022a, 2022c, 2022b, 2023a) regard the metric space or metric strategy as meta knowledge and perform the novel class prediction in a nearest-prototype manner. The model-based methods (Hou et al. 2019; Zhmoginov, Sandler, and Vladymyrov 2022; Li et al. 2019) regard a black-box model as meta knowledge, which leverages it and few data to directly predict model weights or test sample labels. The gradient-based methods (Rusu et al. 2018; Lee et al. 2019; Rajeswaran et al. 2019; Nichol, Achiam, and Schulman 2018; Von Oswald et al. 2021; Raghu et al. 2020; Zhang et al. 2023b) regard the gradient-based optimization algorithm as meta knowledge, which learn to model its hyperparameters (i.e., learning rate (Baik et al. 2020), loss function (Baik et al. 2021), initialization (Finn et al. 2017), preconditioner (Kang et al. 2023), or updata rules (Deleu et al. 2022; Ravi and Larochelle 2017)) such that the base learner can be quickly learned with few gradient updates.

We focus on the gradient-based meta-learning due its good generalization. However, different from existing methods, we presents a new diffusion perspective to model meta-optimizer, which does not need to differentiate through inner-loop path such that the issues of memory burdens and vanishing gradients can be alleviated for improving FSL.

Diffusion model (Ho, Jain, and Abbeel 2020; Nichol and Dhariwal 2021) is a popular type of deep generative models, which models and learns the generation process of target data from random Gaussion noises in a forward diffusion and reverse denoising manner. With the superior properties of diffusion models, the diffusion models have been widely exploited on various vision (Lugmayr et al. 2022) and multi-modal tasks (Rombach et al. 2022; Kawar et al. 2023; Kumari et al. 2023) and achieved remarkable performance improvement. However, there is very few works (i.e., (Roy et al. 2022) and (Hu et al. 2023)) to explore diffusion models for FSL. Specifically, in (Roy et al. 2022), Roy et al. introduce class names as priors and then leverage it and a text2image diffusion model to generate more images for alleviating the data-scarcity issue of FSL. Instead of using text2image diffusion models, Hu et al. (Hu et al. 2023) employ a image2image diffusion model and leverage it to generate more high-similarity pseudo-data for improving FSL.

Different from existing methods that regarding the diffusion model as a component of data augmentation, we find that the gradient descent process is similar to the denoising process, thus we propose to model the gradient descent algorithm as a diffusion model. We note that a concurrent working with our MetaDiff is ProtoDiff (Du et al. 2023). However, different from ProtoDiff that focuses on metric-based meta-learning (i.e., rectifying prototype bias), we target at gradient-based meta-learning, and first reveal the close connections between gradient descent algorithm and diffusion models. Then, a new diffusion-based meta-optimizer is presented for fast adaptation of base-learner.

For a $N$-way $K$-shot FSL problem, it consists of two datasets, i.e., a base class dataset $\mathcal{D}_{base}$ and a novel class dataset $\mathcal{D}_{novel}$. The base class dataset $\mathcal{D}_{base}$ consists of abundant labeled data from base class $\mathcal{C}_{base}$, which is used for assisting the classifier learning of novel classes. The novel class dataset $\mathcal{D}_{novel}$ contains two sub datasets from novel classes $\mathcal{C}_{novel}$, i.e., a training set (call support set $\mathcal{S}$) that consists of $N$ classes and $K$ samples per class and a test set (call query set $\mathcal{Q}$) consisting of unlabeled samples.

Our goal is that leveraging the base class dataset $\mathcal{D}_{base}$ to learn a good meta-optimizer such that the classifier can be quickly learned from few labeled data (i.e., the support set $\mathcal{S}$) to perform the novel class prediction for query set $\mathcal{Q}$.

1) The diffusion process aims to iteratively add a noise from a Gaussian distribution to a target data $x_{0}\sim p_{target}$ to transform $x_{0}$ into $x_{1},x_{2},...,x_{T}$. The final $x_{T}$ tends to become a sample point from the prior distribution $p_{prior}\in\mathcal{N}(\mathbf{0},\mathbf{I})$ when the number of iterations $T$ tends to big enough. The diffusion process aims to learn a noise prediction model $\epsilon_{\theta}(x_{t},t)$ for estimating the added noise at time $t-1$ from $x_{t}$, which is then used to recovery the target data in denoising process. The training object $L$ is as follows:

where $\|\cdot\|_{2}^{2}$ denotes a mean squared error loss. It is worth noting that the above training object defined in Eq. (1) can be performed at any time step $t$ without the iterations of adding noise due to its good closed form at any time step $t$. That is,

where $\beta_{t}\in(0,1)$ is a variance hyperparameter.

2) The denoising process is reverse process of diffusion. Based on the learned noise prediction model $\epsilon_{\theta}(x_{t},t)$, given a start noise $x_{T}\sim p_{prior}$, we can iteratively remove its fraction of noises at each time $t$ and finally recovery the target data $x_{0}$ from the noisy data $x_{T}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$. That is,

where $\sigma_{t}$ is a variance hyperparameter, which is theoretically set to $\sigma^{2}_{t}=\beta_{t}$ in most existing diffusion works (Ho, Jain, and Abbeel 2020; Nichol and Dhariwal 2021).

As introduced in Section 3.2, we can describe the process of denoising and gradient descent process as follows: 1) given a noise data $x_{T}$, the denoising process iteratively performs Eq. (3) to obtain a latent sequence $x_{T-1},x_{T-2},...,x_{0}$. As a result, a target data $x_{0}$ can be recoveried from a Gaussian noise $x_{T}$; and 2) given a randomly initial wight $w_{0}$, the gradient descent process of GDA iteratively performs Eq. (4) to a latent sequence $w_{1},w_{2},...,w_{T}$. As a result, an optimizal weight $w_{T}$ can be obtained from a random weight $w_{0}$. We can see that the denoising process of diffusion models and gradient descent process of GDA are very similar in workflow. This inspires us to think about what is the close connection between the two methods in theory. Let’s take a closer look on Eqs. (3) and  (4), Eq. (3) first can be simplifed as:

Let $\gamma$ denotes the Term 1, $\eta$ be the Term 2, $\xi$ be the Term 3 of Eq. (5), respectively. The Eq. (5) can be simplifed as:

Due to $\gamma>1$ ($\alpha_{t}<1$), we can transform Eq. (6) as follows:

where $Term1$ is denoising term, $Term2$ denotes a momentum update term with hyperparameter $\gamma-1$ (also called exponentially weighted moving average), and $Term3$ is a uncertain term. Comparing Eqs. (4) and (7), we can see that the gradient decent process defined in Eq. (4) is equivalent to the $Term1$ of Eq. (7), which means that Eq. (4) is a special case of denoising process described in Eq. (7) when the $\gamma$ is set to one (i.e., $\gamma=1$), the $\eta$ is regarded as a hyperparameter, and the $\xi$ is set to zero. In particular, it is worth noting that the predicted variable of noise prediction model $\epsilon_{\theta}(x_{t},t)$ is actually the gradient (i.e., $\nabla L(w)$) of model weights. In other words, the denosing process defined in Eq.(7) can be viewed as a generalized and learnable gradient descent algorithm defined in Eq.(4), i.e., a learnable gradent descent algorithm with weight momentum updates and uncertainty estimation where $\gamma$ controls the weight of momentum updates, $\eta$ is a learning rate, and $\xi$ is the degree of uncertainty.

Why set parameters ($\gamma$, $\eta$, and $\xi$) by following Eq. 5? Instead of using manual setting or model learning manner like existing meta-optimizers to set hyperparameters $\gamma$, $\eta$, and $\xi$, respectively, the diffusion models unify the parameter settings by theoretical deduction, i.e., $\gamma=\frac{1}{\sqrt{\alpha_{t}}}$, $\eta=\frac{\beta_{t}}{\sqrt{\alpha_{t}}\sqrt{(1-\overline{\alpha}_{t})}}$, and $\xi=\sigma_{t}$ where $\alpha_{t}=1-\beta_{t}$, $\sigma^{2}_{t}=\beta_{t}$, and $\beta_{t}$ is experimentally set in linear decreasing manner from a small value (e.g., $10^{-4}$) to a large value (e.g., 0.02). The goal of such setting is to ensure that denoising and diffusion processes have approximately the same functional form and the efficiency and robustness of diffusion training (i.e., the training objective defined in Eq. (1) can be performed at any time step $t$ without the iterations from $t=0$ to $t$).

Inspired by the above analysis, we find that the diffusion model is a generalized and learnable form of GDA and its hyperparameter setting have rigorous theoretical derivation, which enables its inspiring advantage (i.e., generation robustness and training efficiency). Based on this, we attempt to leverage a diffusion model to model GDA and then present a new meta-optimizer, i.e., MetaDiff, for fast adaptation of base-learner. It does not need to differentiate through inner-loop path, such that the memory burden and risk of vanishing gradients can be alleviated for improving FSL.

Specifically, given a $N$-way $K$-shot FSL task, we first leverage the embedding network $f_{\varphi}(\cdot)$ to encode the feature $f_{\varphi}(u_{i})$ for each support/query image $u_{i}\in\mathcal{S}\cup\mathcal{Q}$. Then, we randomly initialize a weight $w_{T}\sim\mathbb{N}(\mathbf{0},\mathbf{I})$ for the base learner $g_{w}(\cdot)$, and design a task-conditional UNet (i.e., the noise prediction model $\epsilon_{\theta}(\cdot)$) that regards the features and labels of all support sample $(u_{i},y_{i})\in\mathcal{S}$ as task condition, to estimate the noise to be removed at time $t$. After that, we take the weight $w_{T}$ as the denoising variable and iteratively perform the denoising process from $t=T$ to $t=1$, that is,

Note that we remove the uncertainty term (i.e., $\sigma_{t}z$) for deterministic estimation during inference. After iteratively perform $T$ step, the target weight $w_{0}$ can be obtained as the optimal weight $w$ for base learner $g_{w}(\cdot)$. Finally, we perform class prediction of each query image $u_{i}\in\mathcal{Q}$ by using the learned optimal base learner $g_{w}(\cdot)$. That is,

Here, we only introduce the inference workflow of our MetaDiff-based FSL framework, which is summaried in Algorithm 2. Next, we will introduce the design details of our key component, i.e., the task-conditional UNet $\epsilon_{\theta}(\cdot)$.

Based on this find, as shown in Figure 3, we design a task-conditional UNet from the perspective of gradient estimation as the noise prediction model $\epsilon_{\theta}(\cdot)$. The key idea is predicting noise from the view of gradient estimation instead of a general black-box manner like (Rombach et al. 2022). Specifically, given a base learner weight $w_{t}$ at time $t$ and all features and labels of support samples $(u_{i},y_{i})\in\mathcal{S}$, we first leverage the base learner $g_{w_{t}}(\cdot)$ with model weights $w_{t}$ to compute the loss $L_{w_{t}}(\mathcal{S})$ of of all support samples. That is,

where $|\cdot|$ is the number of support samples and $loss\_fun(\cdot)$ is a loss function. Instead of cross-entropy loss, we employ a simple L2 loss to implement the $loss\_fun(\cdot)$, that is,

where $w_{t,y_{i}}\in w_{t}$ is the class prototype of label $y_{i}$. The intuition of such design is moving the class prototype towards the center of all labeled sample for each class $y_{i}$, which is more matching for the rationale of prototype classifier. Then, the gradient $\nabla L_{w_{t}}(\mathcal{S})$ regarding weights $w_{t}$ can be obtained as the initial noise estimation for the base learner $g_{w_{t}}(\cdot)$ at time $t$. To obtain more accuracy noise estimation, we take the initial noise estimation $\nabla L_{w_{t}}(\mathcal{S})$ as inputs and then design a conditional UNet fusing time embedding $t$ to predict the noise to be remove at time $t$ for base learner $g_{w}(\cdot)$.

As shown in Figure 3, the UNet consists of two encoder blocks (EB), a bottle block (BB) and two decoder blocks (DB). At each encoder step, we halve the number of input features and then remain unchanged at bottle step, but the number of features is doubled at each decoder step. The details of each encoder, bottle, decoder block are all similar, which contains a feature transform layer, a time embedding layer, and a ReLU activation layer. Note that we remove the ReLU activation layer in the final decoder block for estimating gradients. At each block, its output is obtained by first feeding the output of previous block and time step $t$ into the feature transform and time embedding layers, respectively, and then fusing them in an element-by-element product manner, finally followed by a softplus activation.

To this end, we follow episodic training strategy (Vinyals et al. 2016) and construct a large number of $N$-way $K$-shot tasks from base class dataset $\mathcal{D}_{base}$. Then, given a constructed $N$-way $K$-shot tasks $\tau$, based on its origin label $k^{\prime}$ of each class $k=0,1,..,N-1$ in the base classes $\mathcal{C}_{base}$, we extract all samples that belongs to the origin label $k^{\prime}$ of each class $k=0,1,..,N-1$ from the base class dataset $\mathcal{D}_{base}$, as the auxiliary dataset $\mathcal{D}_{base}^{\tau}$. The labeled data is very sufficient in the auxiliary dataset $\mathcal{D}_{base}^{\tau}$ because it contains all labeled data belonging to class $k^{\prime}$ in $\mathcal{D}_{base}$, thus we can leverage it to learn a base learner $g_{w}(\cdot)$ such that the target weight $w_{0}$ can be obtained for each task $\tau$. Finally, we leverage the target weight $w_{0}$ of all constructed tasks to train our MetaDiff meta-optimizer in a diffusion manner. That is,

During training, our MetaDiff does not require backpropagation along the inner-loop optimization path and calculating second-order derivatives for learning meta-optimizer such that the memory overhead and the risk of vanishing gradient can be effectively alleviated for improving FSL. The complete diffusion procedure is summaried in Algorithm 1.

Our MetaDiff falls into the type of gradient-based meta-learning, thus we mainly select various state-of-the-art gradient-based meta learning methods as our baselines. We evaluate our MetaDiff and these baselines on Imagenet derivatives. The experimental results are shown in Table 1. Among them, iMAML(Rajeswaran et al. 2019), MAML (Finn et al. 2017), ALFA (Baik et al. 2020), ANIL (Raghu et al. 2020), COMLN (Deleu et al. 2022), GAP (Kang et al. 2023), LEO (Deleu et al. 2022), and ClassifierBaseline (Chen et al. 2021), are our key competitors, which also focus on learning GDA but modeling its hyperparameters.

Table 1 shows the results of various gradient-based meta-learning methods on miniImagenet and tieredImagenet. From these results, we find that (i) our MetaDiff achieves superior or comparable performance on all tasks, which exceeds most state-of-the-art gradient-based meta-learning by around 1% $\sim$ 3%. This verifies the effectiveness of our MetaDiff; and (ii) Our MetaDiff achieves consistent improvement on Conv4 and ResNet12 backbones for all tasks, which is reasonable because our MetaDiff mainly focuses on the adaptation of classification head. This also verifies the universality of our MetaDiff on various backbones.