Cooperative Meta-Learning with Gradient Augmentation

In recent, meta-learning successfully covers a diverse application [Hospedales et al., 2021]. Gradient-based meta-learning optimizes a bilevel optimization problem [Colson et al., 2007] where it has a task-adaptation (inner loop) learning a new task with a few shot samples from meta-initialization parameters and meta-optimization (outer loop) finding proper meta-initialization parameters from an inner loop update. Many variants of MAML have been studied in various domains [Yin et al., 2020, Obamuyide and Vlachos, 2019, Collins et al., 2022, Lee et al., 2021]. BOIL [Oh et al., 2020] tackles the feature reuse problem in meta-optimization and freezes the classifier in task-adaptation. Sharp-MAML [Abbas et al., 2022] leverages sharpness-aware minimization to solve a bilevel optimization problem. In this work, we propose a new meta-learning framework that can be applied to any gradient-based meta-learning.

While maintaining the exact computational graph for inference, many works to boost the performance of the model have been studied. Auxiliary training adds auxiliary classifiers connected in intermediate layers [Szegedy et al., 2015, Zhang et al., 2020] and multi-task learning simultaneously learns multiple related tasks and the knowledge from multi-task can be reused by the others [Yang and Hospedales, 2017]. Unlike previous methods, multi-branch frameworks [Kim et al., 2021, Xie and Du, 2022, Liang et al., 2022] shared intermediate layers and split multi-branch under the same task which utilizes knowledge distillation [Hinton et al., 2015] transferring the knowledge to enhance independent branches. Zhu et al. [2018] split the model into several sub-networks and made an ensemble logit to teach individual sub-networks. Song and Chai [2018] introduces multiple heads from the same network to improve the generalization of the model.

To improve the generalization performance, various ways to impose constraints on model structure and gradients by noise have been studied. Hinton and Roweis [2002] uses gaussian gradient noise schedule to train the embedding model. Dropout [Srivastava et al., 2014] randomly drops the connections during training which introduces the random noise into forward propagation. Similarly, Huang et al. [2016] randomly disconnects the layers during training. Neelakantan et al. [2015] shows injecting noise to gradient works very deep architecture. GradAug [Yang et al., 2020] generates meaningful noise in gradients rather than random noises by multiple data augmentation and pruning of the model.

Cooperative meta-learning leverages the advantages of both multi-branch framework and regularization by noise in the gradient-based meta-learning domain. CML only introduces a co-learner and trains it in meta-optimization to augment a meta-gradient by sharing the feature extractor. It induces a regularization effect by injecting noise into the meta-gradient without multiple forwarding or making a sub-network such as pruning.

In this work, we divide the DNN model used for meta-learning into two groups: the feature extractor $\psi$ which extracts the features, i.e., body layers of DNN and the meta-learner $\mathcal{\theta}$ predicting outputs based on the features, i.e., classifier. We sample a set of tasks $\{\mathcal{T}\}^{N}_{i}$ containing N tasks from the task distribution $p(\mathcal{T})$. DNN model represented by a $f_{(\psi,\theta)}$ is trained using samples from each task $\mathcal{T}_{i}$ under the two optimization loops. These samples $\mathcal{D}_{i}$ are divided into support set $\mathcal{D}^{S}_{i}$ and query set $\mathcal{D}^{Q}_{i}$ which are used in the inner loop and outer loop, respectively. MAML consisting of two optimization loops which are the inner loop and outer loop tries to find well-generalized meta-initialization parameters during meta-training. In the inner loop, we update task-specific parameters from meta-initialization parameters $(\psi,\theta)$ using the support set with an outer step size of $\alpha$.

and takes totally $M$-updates for task-specific parameters.

We will consider one gradient step for the rest for simplification. After task-adaptation in the inner loop, we compute each task loss for the query set with task-specific parameters $(\mathcal{\psi}^{\prime}_{i},\mathcal{\theta}^{\prime}_{i})$. By summing all task losses, meta-optimization optimizes the following objectiveness:

In the outer loop, we update meta-initialization parameters with $N$ task losses using meta-gradient by meta-optimization with an outer step size of $\beta$.

In meta-testing, we verify the trained meta-initialization parameters. The inner loop adapts to the new task with a support set that remains the same as in meta-training. However, the outer loop only computes the accuracy using a query set for each task. There is no meta-optimization process in the outer loop of meta-testing.

Our proposed framework includes an additional module called co-learner $\phi$ inducing gradient augmentation in the meta-optimization. The co-learner can consist of a convolution layer or a fully connected layer, depending on the task. In meta-training, CML performs task-adaptation with the feature extractor and meta-learner in the inner loop same as the original gradient-based meta-learning such as MAML. The co-learner is added to the feature extractor during meta-optimization in the outer loop. Note that the co-learner only intervenes in the outer loop to perform meta-optimization with the feature extractor and meta-learner. In other words, the co-learner does not perform task-adaptation for the current task in the inner loop, therefore, it has implicit knowledge of tasks in the previous sampled batch. As a result, it has a different representation from the naive meta-learner. In this framework, the meta-learner and co-learner always share a representation of the feature extractor. Hence, their gradients are aggregated in the feature extractor, resulting in gradient augmentation.

Formally, we sample a support set $\mathcal{D}_{i}^{S}$ and query set $\mathcal{D}_{i}^{Q}$ from a new task $\mathcal{T}_{i}$. Also, we denote the initial parameters for the feature extractor, meta-learner and co-learner as $\mathcal{\psi}$, $\mathcal{\theta}$ and $\mathcal{\phi}$ and two models in our framework: model $f^{m}$, which composes of a shared feature extractor and a meta-learner, and $f^{c}$, which composes of a shared feature extractor and a co-learner. In the inner loop, the feature extractor and the meta-learner update the parameters $\psi_{i}^{\prime}$ and $\mathcal{\theta}_{i}^{\prime}$ with $\mathnormal{M}$-updates from a batch of $\mathcal{D}_{i}^{S}$, respectively. However, the co-learner does not update the parameters $\mathcal{\phi}_{i}^{\prime}$ in the inner loop. Therefore, task-specific parameters $\mathcal{\psi}_{i}^{\prime}$, $\mathcal{\theta}_{i}^{\prime}$ and $\mathcal{\phi}_{i}^{\prime}$ are as follows:

where $\alpha$ is an inner step size which is a fixed hyper parameters. Unlike $\psi_{i}^{\prime}$ and $\theta_{i}^{\prime}$ which are updated for the current task $\mathcal{T}_{i}$ in the inner loop, $\mathcal{\phi}_{i}^{\prime}$ keeps the updated parameters from the previously sampled tasks in the outer loop. In the outer loop, our model $f$ updates meta-initialization parameters from $\mathcal{D}_{i}^{Q}$ with task-specific parameters updated by $\mathcal{D}_{i}^{S}$. Our purpose is to converge to $\psi$, $\theta$ and $\phi$ that minimize Eq.(6) with the loss of the meta-learner and co-learner.

where $\gamma$ is the loss scaling factor. The feature extractor and meta-learner have task-specific parameters $\mathcal{\psi}_{i}^{\prime}$ and $\mathcal{\theta}_{i}^{\prime}$ with knowledge about the current task $\mathcal{T}_{i}$, but the co-learner has the parameters $\mathcal{\phi}$ that have been updated by meta-optimization on the previous sampled N tasks.

Then, we compute the meta-gradient with $N$ task losses for the query set $\mathcal{D}_{i}^{Q}$. It is created by gradient augmentation, where the gradient noise from the co-learner is added to the existing gradient. From Eq.(7), we update the meta-initialization parameters with an outer step size of $\beta$ by meta-optimization in the outer loop. The updated $\mathcal{\psi},\mathcal{\theta}$ and $\phi$ are initialized with meta-initialization parameters for meta-testing. Lastly, our framework can infer with CML and CML^† using meta-testing dataset $D_{test}=(D_{test}^{S},D_{test}^{Q})$ in meta-testing phase, and CML performs the task-adaptation as follows:

The model then evaluates against $\mathcal{D}^{Q}_{test}$ by using the adapted parameters $\mathcal{\psi^{\prime}}$ and $\mathcal{\theta^{\prime}}$ like standard MAML. Therefore, it does not require any additional inference cost for the co-learner.

On the other hand, in Eq.(9), CML^† performs the task-adaptation only for $\mathcal{\psi}$. Note that $\phi$ of the co-learner does not perform task-adaptation and has existing meta-initialization parameters. Then we evaluate model with $\mathcal{\psi}^{\prime}$ and $\phi$ against $\mathcal{D}^{Q}_{test}$. Our CML algorithm is shown in Algorithm 1.

Next, we demonstrate that the gradient calculated from the co-learner converges theoretically when it is combined into meta-gradients. Our meta-gradient is updated by combining the gradients from the meta-learner($\theta^{\prime}$) and co-learner($\phi$) in the feature extractor($\psi^{\prime}$). We represent the loss function of the base network with $\psi^{\prime}$, $\theta^{\prime}$ to be $\mathcal{L}(\psi^{\prime},\theta^{\prime})$ after the task-adaptation.

We evaluate the performance of CML with MAML as a baseline in K-shot sinusoidal regression. The amplitude and phase of the sinusoidal wave follow the ranges of [0.1,5.0] and [0,$\pi$]. Each task consists of datapoints $\mathbf{x}$, $\mathbf{y}$ of a sinusoidal wave. The input $\mathbf{x}$ is uniformly sampled in the range [-5.0,5.0]. The loss function for comparing predicted $\mathbf{y}$ and target $\mathbf{y}$ uses mean-squared error. The baseline consists of 2 hidden layers of size 40 with ReLU nonlinearities, 1 input layer and 1 output layer following [Finn et al., 2017]. For CML, the regressor is additionally attached with 1 hidden layer of size 40 with ReLU nonlinearities and 1 output layer as a co-learner. In meta-training, we use K $\in$ $\{$5,10,20$\}$ samples as training examples and train using a batch size of 4, one inner-gradient step, a fixed step size of 0.01 and our loss scaling factor $\gamma$ of 0.2. For meta-testing, we evaluate adaptation with one gradient step for K=5, 10, and 20 test points. Each model predicts the target sinusoidal wave through the given K test points. Furthermore, the co-learner is deleted and it is only evaluated from the feature extractor and meta-learner as the original model. Figure 2 shows that our CML performs better than MAML for 5, 10, and 20 shots. It means that our framework adapts well to simple networks and shows better generalization performance than the original framework.

We compare the performance of the proposed method on few-shot image classification with several gradient-based meta-learning algorithms including MAML [Finn et al., 2017], MAML++ [Antoniou et al., 2018], BOIL [Oh et al., 2020] and Sharp-MAML [Abbas et al., 2022]. In this experiment, we evaluate the performance of 5-way 1/5-shot problems on MiniImagenet datasets. In CML, the co-learner uses two convolution layers and a fully connected layer. Our loss scaling factor $\gamma$ is fixed at 0.5 for all methods. We also evaluate the performance of the co-learner.

In this experiment, we evaluate CML on a few-shot node classification of graph neural networks (GNNs). Few-shot node classification aims to achieve fast adaptation to new node tasks that are unseen during training. We also define an N-way K-shot problem and consider node tasks $\mathcal{T}_{node}$ which consist of support nodes $\mathcal{D}_{node}^{\mathcal{S}}$ and query nodes $\mathcal{D}_{node}^{\mathcal{Q}}$. For performance comparison, we use G-Meta [Huang and Zitnik, 2020] and AMM-GNN [Wang et al., 2020], which belong to gradient-based meta learning with GNN, as a baseline and evaluate on the CoraFull, Amazon-Computer and CiteSeer datasets [Sen et al., 2008, Shchur et al., 2018]. We also perform 5-way 1/5-shot and 2-way 1/5-shot, respectively. Our base model follows [Tan et al., 2022], using GCN as an encoder of hidden size 16 and a fully connected layer as a meta-learner. We train using the Adam optimizer for a step size of 0.001. Also, we set the inner gradient-steps of 20 with a step size of 0.05. In CML framework, our co-learner additionally includes 1 hidden layer of size 16 and a fully connected layer as the output layer using the loss scaling factor $\gamma$ of 0.2. As shown in Table 3, our framework outperforms the baseline method on node classification. In this experiment, the co-learner improves the performance of the encoder and meta-learner, despite having a very simple network structure. It shows that our framework is suitable for solving the few-shot problems and is applicable to various DNN methods related to meta-learning.

In this section, we analyze the effect of gradient augmentation by a co-learner. All of our experiments are evaluated on MiniImagenet 5-way 5-shot, and the network structure and experimental settings are as in Section 4.2.

In this section, we conduct an experiment to justify the validity of our framework’s structure. Our framework requires more parameters during meta-training due to the addition of the co-learner. Therefore, we compare the parameter sizes of CML, CL and MAML with more parameters, called More-MAML, to demonstrate that our framework does not simply improve performance by having more parameters. In this experiment, CML and CL follow the same network architecture as Section 4.2, while More-MAML has additional convolution layers. From Table 4, we can see that More-MAML, CL and CML have 232K, 203K and 203K parameters, respectively, during meta-training. Note that the CML has fewer parameter sizes than More-MAML and CL in meta-testing, but shows better performance on MiniImagenet datasets. Also in (1), CL_⋆ shows a performance degradation when inferring with a meta-learner like CML. In setting (2), both CL^† and CML^† use the co-learner to evaluate performance. We observe that without performing adaptation, CL^† leads to a deterioration in the model’s inferential capabilities. It emphasizes that adaptation is essential in the general case, and that our approach has a uniquely structured framework. Notably, although our co-learner does not perform task-adaptation during meta-testing in CML^†, it outperforms models with a similar number of parameters while achieving the highest accuracy. In this experiment, our findings highlight that having more parameters in meta-learning does not necessarily lead to improved performance, while our framework demonstrates an effective learning framework to address this limitation.