Stochastic Deep Networks with Linear Competing Units for Model-Agnostic Meta-Learning

1 INTRODUCTION . When we train machine learning models on problems with limited amounts of training data , we can not usually get good predictive performance ( Lai , 2019 ; Sculley et al. , 2015 ) . This comes in contrast to the human ability to quickly derive information from a range of different tasks , and then adapt to a new task with limited available new examples ( Kühl et al. , 2020 ; Castro et al. , 2008 ) . In essence , this capability of the mind of learning how to learn ( Black et al. , 2006 ) has inspired researchers to investigate the concept of Meta-Learning ( ML ) ( Lake et al. , 2017 ; Vilalta & Drissi , 2002 ; Wang et al. , 2016 ; Zoph & Le , 2017 ; Flennerhag et al. , 2020 ; Lake et al. , 2015 ; HutsebautBuysse et al. , 2019 ) . There is a large variety of deep learning methods for ML ( Finn et al. , 2019 ; Baik et al. , 2020 ; Andrychowicz et al. , 2016 ) . Specifically , Finn et al . ( 2017 ) presented the Model-Agnostic MetaLearning ( MAML ) algorithm that enables tuning the parameters of a trained network to quickly learn a new task with only a few gradient updates . To get away with the entailed second-order computations , which are expensive , several researchers proposed appropriate first-order approximations for MAML ; such works are the First-Order MAML ( FOMAML ) of Finn et al . ( 2017 ) and the Reptile algorithm of Nichol et al . ( 2018 ) . Recently , several researchers have also considered Bayesian inference-driven methods for deep learning ML , extending upon older works built for conventional machine learning approaches ( Yoon et al. , 2018 ; Finn et al. , 2018 ; Ravi & Beatson , 2019 ; Patacchiola et al. , 2020 ; Zou & Lu , 2020 ; Grant et al. , 2018 ; Chen et al. , 2020 ) . This work proposes a different regard toward improving generalization capacity for deep networks in the context of ML . Specifically , our proposed approach relies on the following main concepts : • The concept of sparse learned representations . For the first time in the literature of deep network-driven ML , we employ a mechanism that inherently learns to extract sparse data representations . This consists in replacing standard unit nonlinearities ( e.g. , ReLU ) with a unit competition mechanism . Specifically , ( linear ) units are organized into blocks . Presented with some input , the units within a block engage in a competition process with only one winner . The outputs of all units except for the winner are zeroed out ; the output of the winner retains its computed value ( local winner takes-all , LWTA , architecture ) . • The concept of stochastic representations . We establish a stochastic formulation for the previously described competition process . Specifically , we postulate that , within a block of competing units , winner is selected via sampling from an appropriate Categorical posterior . The corresponding winning probability of each unit is proportional to its linear computation ( thus depending on the layer input ) . Via this competition process , we yield stochastic representations from network layers , that is representations that may change each time we present to the network layer exactly the same input . Based on the results from existing approaches , we posit that the proposed treatment of the ML problem , which combines learned representation sparsity and stochasticity , will be extremely beneficial to the deep learning community . We dub our approach Stochastic LWTA for ML ( StochLWTA-ML ) . We perform a variational Bayes treatment of the proposed model . We opt for a full Bayesian treatment , by also handling network weights as latent variables . That is , we elect to impose an appropriate prior over the network weights and fit approximate ( variational ) posteriors . We evaluate our approach on a number of standard benchmarks in the field , namely Omniglot ( Lake et al. , 2017 ) , Mini-Imagenet ( Vinyals et al. , 2016 ) and CIFAR-100 ( Krizhevsky , 2009 ) . We show that our approach offers a variety of advantages over the current state-of-the-art methods , namely : ( i ) incurring reduced predictive error rate compared to the currently state-of-the-art methods in the field ; ( ii ) obtaining this performance with networks that comprise one order of magnitude less trainable parameters , and therefore give rise to better computational efficiency and imposed memory footprint . The remainder of this paper is organized as follows : In Section 2 , we briefly review related work . Section 3 introduces our approach and provides the related training and prediction algorithms . In Section 4 , we perform a thorough experimental evaluation of StochLWTA-ML , and compare our findings to the current state-of-the-art . In the final Section 5 , we end up with the conclusions of our work , and suggest lines of further research . 2 RELATED WORK . 2.1 LWTA LAYERS IN DEEP LEARNING . LWTA layers are not new in the field of deep learning ; see , e.g. , Srivastava et al . ( 2013 ) . Although not much work has been pursued along these lines , the recent works of Panousis et al . ( 2019 ; 2021 ) and Voskou et al . ( 2021 ) have spurred some fresh interest in the field . These works have presented alternative implementations of the basic concepts of LWTA units in the context of diverse deep network architectures . Specifically , Panousis et al . ( 2019 ) propose a stochastic LWTA formulation which is founded upon the Indian Buffet process ( IBP ) prior , borrowed from nonparametric statistics ; they use this architecture to effect data-driven network compression . In their follow-up paper ( Panousis et al. , 2021 ) , they exploit the same technique to train adversarially-robust deep networks . On the other hand , Voskou et al . ( 2021 ) consider a different incarnation of stochastic LWTA architectures , which relies on sampling the winner from a Categorical posterior , driven from the layer input . This layer architecture is used to replace dense ReLU layers in Transformer networks ; it is then shown to yield important benefits in a Sign-Language Translation benchmark . This paper is different from the previous works in various ways : ( i ) stochasticity does not stem from utilization of the IBP ; we rather adopt an approach similar to Voskou et al . ( 2021 ) ; ( ii ) we do not use the proposed architecture as a replacement for a specific type of layer in a greater network architecture ( Transformer ) that remains largely unchanged ; instead , we build completely new networks using these layers ; ( iii ) we perform a full Bayesian treatment , by treating network weights as random variables ; we do not employ this construction as a means of compressing the weights at prediction time , contrary to Panousis et al . ( 2019 ) , Voskou et al . ( 2021 ) ; instead , we perform weight sampling at prediction time as a means of improving accuracy ; and ( iv ) for the first time , we examine how these principles perform in the context of deep network-driven ML . Note that , apart from LWTA architectures , other data-driven sparsity models have also been proposed recently , e.g . Lee et al . ( 2018 ) , Kessler et al . ( 2021 ) . However , none of these have been developed or evaluated in the context of an ML setting . 2.2 MODEL-AGNOSTIC META-LEARNING . ML ( Schmidhuber , 1987 ) , also referred to as learning to learn ( Thrun & Pratt , 1998 ) , has the goal of obtaining machine learning models that can learn new tasks with limited availability of data , and via only a few gradient steps , by capitalizing upon information stemming from previously learned tasks . In this context , Finn et al . ( 2017 ) suggested a model-agnostic algorithm for ML , that can be applied to any model trained via gradient descent . The introduced MAML algorithm initializes model parameters in a way that can be quickly adapted to several types of new tasks . Let us consider a model with parameters θ and a parametric form fθ . When the model is adapting to an unseen task Ti , sampled from the distribution over tasks P ( T ) , MAML initially runs few steps of inner-loop gradient descent that yields the task-specific parameter set θ′i = θ − α∇θLTi ( fθ ) ( 1 ) where α is the step size hyperparameter and LTi denotes the loss on the task Ti . Subsequently , training proceeds to optimize the function fθ′i with respect to the model parameters θ . Assuming that the batch of tasks has size M , we can define the targeted meta-objective as : Lmeta ( θ ) = M∑ i=1 LTi ( fθ′i ) = M∑ i=1 LTi ( fθ−α∇θLTi ( fθ ) ) . ( 2 ) Optimization of this meta-objective over θ yields the outer-loop update : θ ← θ − β∇θ M∑ i=1 LTi ( fθ′i ) ( 3 ) where β stands for the outer-loop learning rate . Finally , as MAML involves expensive computations stemming from the second-order updates of Eqs . ( 2 ) and ( 3 ) , Finn et al . ( 2017 ) developed a first-order approximation that reduces the outer-loop update ( 3 ) to : θ ← θ − β M∑ i=1 ∇θ′iLTi ( fθ′i ) ( 4 ) In other words , FOMAML computes the gradients with respect to the updated parameter values θ′i , but omits the gradients of θ′i with respect to θ . In a different fashion , the Reptile algorithm of Nichol et al . ( 2018 ) disposes the costly secondorder updates of MAML by : ( i ) applying some inner-loop gradient descent steps using ( 1 ) ; and ( ii ) suggesting a new outer-loop that simply subtracts the parameters , θ , from the updates θ′i ( instead of computing derivatives ) : θ ← θ + β ( θ′i − θ ) . ( 5 ) 3 PROPOSED APPROACH . 3.1 ARCHITECTURE . Let us denote as x ∈ RI an input vector presented to a dense ReLU layer of a conventional deep neural network , with corresponding weights matrix W ∈ RI×O . The output of the layer is the vector y ∈ RO and is fed to the subsequent layer . In our approach , a ReLU unit is replaced by J competing linear units , organized in one block ; in the following , we denote with R the number of blocks in a layer . The input x is now presented to each block through weights that are organized into a three-dimensional matrix W ∈ RI×R×J . Then , the j-th competing unit within r-th block computes the sum ∑I i=1 ( wi , r , j ) · xi . Competition means that , of the J units in the block , one unit ( the ” winner ” ) will present its linear computation to the next layer ; the rest will present zero values . Traditionally in the literature , the winner unit is selected to be the unit with greatest linear computation . Recently , stochastic competition principles have been considered , e.g . Panousis et al . ( 2019 ; 2021 ) , Voskou et al . ( 2021 ) . Let us denote as y ∈ RR·J the output of an LWTA layer ; this is composed of R subvectors yr ∈ RJ and is sparse , since all units except for one , in each block , yield zero values . Let us introduce the discrete latent indicator vector ξ ∈ one hot ( J ) R to denote the winner units in the R blocks that constitute a considered stochastic LWTA layer . This vector comprises R component subvectors , where each component entails one non-zero value at the index position that corresponds to the winner unit of the respective LWTA block . On this basis , the output y of the stochastic LWTA layer ’ s ( r · j ) -th component yr , j is defined as : yr , j = ξr , j I∑ i=1 ( wi , r , j ) · xi ∈ R ( 6 ) where we denote as ξr , j the j-th component of ξr , and ξr ∈ one hot ( J ) holds the r-th subvector of ξ . In Eq . ( 6 ) , we postulate that the latent winner indicator variables are drawn from a Categorical distribution which is proportional to the intermediate linear computation that each unit performs . Therefore , the stronger the linearity the higher the chance of the unit winning the stochastic competition within its block . In detail , we postulate that , a posteriori , the winner distributions yield : q ( ξr ) = Categorical ( ξr|softmax ( I∑ i=1 [ wi , r , j ] J j=1 · xi ) ) ( 7 ) where [ wi , r , j ] Jj=1 denotes the vector concatenation of the set { wi , r , j } Jj=1 . A graphical illustration of the proposed stochastic architecture is provided in Fig . 1 . As a network composed of such ( StochLWTA ) layers entails latent variables ξ , we need to perform a Bayesian network treatment to perform effective parameter training . We opt for a ( approximate ) stochastic gradient variational Bayes treatment ( Kingma & Welling , 2014 ) , for scalability purposes . This means that the used objective function takes the form of an evidence lower-bound ( ELBO ) objective , as we describe next . In our work , we take one step further : we also elect to infer a posterior density over the network weightsW , as opposed to obtaining point-estimates . This results in a second source of stochasticity for our approach , which may further increase its generalization capacity under uncertain conditions arising from limited training data availability . Note that the use of these posteriors is totally different from Panousis et al . ( 2019 ) and Voskou et al . ( 2021 ) : therein , posterior variance is used for compressing posterior mean bit-precision ; then , predictions are performed using only the compressed posterior mean . Instead , in our work we sample multiple times from the trained weight posterior and perform model averaging ( in a Bayesian sense ) , as a means of increasing generalization capacity . We postulate : q ( vec ( W ) ) = N ( vec ( W ) |µ , diag ( σ2 ) ) ( 8 ) where ψ ≜ { µ , σ2 } are the means and variances of the Gaussian weight posteriors , respectively .

Summary: The paper proposes to use neural networks with so-called "stochastic local winner-takes-all (LWTA)" activations in the place of standard ReLU activations. These stochastic LWTA activations result in a model with sparse representations. The sparse gating is treated as a random variable. The authors also treat the weights of the neural network as random variables. Both sets of random variables are learnt using variational inference.