B-SMALL: A Bayesian Neural Network Approach to Sparse Model-Agnostic Meta-learning

The ability to adapt and learn new models with small amounts of data is a critical aspect of several systems such as IOTs, secure communication networks, biomedical signal processing, image processing etc. Traditional signal processing has addressed such problems using Bayesian and sparse signal processing techniques under a model-driven approach, incorporating several statistical assumptions on the distribution of input data. However, the modern era of artificial intelligence, brings in the promise of model-free processing using various machine learning algorithms, with no assumptions required on the statistical properties of the signals involved.

Among several machine learning approaches proposed to deal with low-data regimes [1, 2], meta-learning is a simple yet efficient technique which aims at obtaining rapid adaptation across various tasks¹¹1Often, task refers to a subset of observations sampled from the original dataset in such a way that only a subset of the final prediction problem can be solved in the task., given small amount of data for updating the parameters pertaining to each task. In particular, model-agnostic meta-learning (MAML) is an algorithm that trains a model’s parameters such that a small number of gradient updates will lead to fast learning on a new task. Specifically, MAML obtains a meta-initialization at meta-training phase using task-specific training, and this initialization is rapidly adapted to a new task by using gradient descent in the meta-test phase. MAML is a baseline for any state-of-the-art few-shot learning method since it has been used for supervised classification, regression and reinforcement learning in the presence of task variability. Furthermore, MAML substantially outperforms techniques that use pre-training as initialization. In order to further improve on the adaptation and accuracy performance of MAML, several authors have proposed modifications such as introducing novel regularizers by analysing the optimization landscape [3], feature reuse perspective based ANIL framework [4], a meta-regularizer using information theoretic approaches for mitigating the memorization problem in MAML [5], etc.

In signal processing based applications, such as distributed signal processing, there is a need for a technique that rapidly adapts in a distributed manner using low amount of heterogeneous data at each sensor node. Furthermore, it is essential that these machine learning models be computationally simple and memory-efficient in terms of the number of parameters they require [6]. The inherent structure of the MAML algorithm lends itself in such scenarios since the task level learning in the inner iteration can be associated to per-node learning, while outer iteration parameter update agglomerates the updates from neighboring nodes, effectively enabling inference capability at each node. However, a challenge in the existing meta-learning approaches is their tendency to overfit, thereby defeating the true purpose of designing such networks [7]. It is well-known that incorporating sparsity constraints during model training guarantees statistical efficiency and robustness to overfitting, hence improving generalization performance on previously unseen tasks [8]. In the context of compact meta-learning, network pruning [9] and regularization [10, 11] have led to sparse meta-learned models without compromising generalization performance. Several methods have been proposed in order to combine deep networks and probabilistic methods for few-shot learning. In particular, in [12], the authors employ hierarchical Bayesian models for few shot learning. In [13], the authors employ a graphical model via a hierarchical Bayesian model that includes prior distribution over the weights and hyperparameters of the meta-learning model.

A popular approach for incorporating uncertainty in deep networks is using the Bayesian neural networks (BNN) [14, 15]. Although exact inference in BNNs is not possible [16], approximations based on backpropagation and sampling have been effective in incorporating uncertainty into the weights [15]. Furthermore, these networks can be made sparse, and eventually compressed to obtain light neural networks [17]. However, so far, conventional BNNs directly learn only the posterior weight distribution for a single task and have not been employed in the meta-learning framework.

We demonstrate the effectiveness of this technique in achieving sparse models with improved accuracy on well-known datasets in the context of classification as well as regression²²2Code for the experiments can be found on github at https://github.com/anishmadan23/B-SMALL. Finally, we present a use-case for the proposed B-SMALL algorithm as distributed algorithms for sensor networks.

In this section, we describe MAML, an optimization based meta-learning paradigm, followed by description of the Bayesian neural network and the SVD paradigm.

Consider a task distribution $p(\mathcal{T})$ over the set of tasks $\mathcal{T}$ that encompasses the data points in $\mathcal{D}=\{x_{j},y_{j}\}$, for $j=1,\ldots,N$. These tasks $\mathcal{T}_{i}\in\mathcal{T}$ are used for meta-training a model $p(y|x,\mathbf{w})$ which predicts $\mathbf{y}$ given inputs $\mathbf{x}$ and parameters $\mathbf{w}$. We adopt the SVD framework to introduce sparsity and reduce overfitting in our meta-learning pipeline, i.e., we maximize the variational lower bound and accordingly modify the loss function of MAML in the inner and outer loop as follows:

Here, similar to the previous section, $\phi$ denotes the variational parameters given by $\phi=[\theta,\alpha]$. Further, $\alpha$ is interpreted as the dropout rate. Note that SVD enables us to have individual dropout rates for each neuron that are learnable. Furthermore, the regularization term is such that $\alpha_{i,j}\rightarrow+\infty$ for several neurons. A high dropout value implies we can effectively ignore the corresponding weight or neuron and remove it from the model, leading to lighter neural network models. Also, $\mathcal{L}^{D}_{T_{i}}(\phi)$ takes the form similar to the cross entropy loss for discrete classification problem, and squared loss in the case of regression problem [17].

In this section, we illustrate the performance of the proposed B-SMALL approach in the context of both, classification and regression. We evaluate the classification performance on few-shot image recognition benchmarks such as the Mini-Imagenet[19] and CIFAR-FS[20] datasets [18]. The setup is a $N$-way, $K$-shot experiment, where we randomly select $N$ classes and choose $K$ images/samples for each class at every training step. All the models in different experiments are trained for $60000$ steps. We measure sparsity as the ratio of total number of weights above a certain threshold $\eta$, and total number of weights. We set $\eta=3$ for all our experiments, and consider those neurons as dropped out when $\log\alpha_{i,j}>\eta$, where $\alpha_{i,j}$ is the variational parameters in (4).

In this paper, we proposed the B-SMALL framework, a sparse BNN-based MAML approach for rapid adaptation to various tasks using small amounts of data. The parameters of the BNN are learnt using a sparse variational loss function. We demonstrated that the proposed framework outperformed MAML in most of the scenarios, while resulting in sparse neural network models. The results obtained builds on the theory that often, in deep learning, we have more parameters as compared to training instances, and such models are prone to overfitting [24]. This gap is amplified in meta-learning since it operates in the low-data-regime and therefore it is important to use regularization technique as in B-SMALL. This helps to reduce the parameter footprint thereby reducing overfitting, and boosts generalization performance. As a future work, we plan to design and analyse B-MAML type algorithms for distributed processing.