Bayesian Context Aggregation for Neural Processes

1 INTRODUCTION . Estimating statistical relationships between physical quantities from measured data is of central importance in all branches of science and engineering and devising powerful regression models for this purpose forms a major field of study in statistics and machine learning . When judging representative power , neural networks ( NNs ) are arguably the most prominent member of the regression toolbox . NNs cope well with large amounts of training data and are computationally efficient at test time . On the downside , standard NN variants do not provide uncertainty estimates over their predictions and tend to overfit on small datasets . Gaussian processes ( GPs ) may be viewed as complementary to NNs as they provide reliable uncertainty estimates but their cubic ( quadratic ) scaling with the number of context data points at training ( test ) time in their basic formulation affects the application on tasks with large amounts of data or on high-dimensional problems . Recently , a lot of interest in the scientific community is drawn to combinations of aspects of NNs and GPs . Indeed , a prominent formulation of probabilistic regression is as a multi-task learning problem formalized in terms of amortized inference in conditional latent variable ( CLV ) models , which results in NN-based architectures which learn a distribution over target functions . Notable variants are given by the Neural Process ( NP ) ( Garnelo et al. , 2018b ) and the work of Gordon et al . ( 2019 ) , which presents a unifying view on a range of related approaches in the language of CLV models . Inspired by this research , we study context aggregation , a central component of such models , and propose a new , fully Bayesian , aggregation mechanism for CLV-based probabilistic regression models . ∗Correspondence to : Michael.Volpp @ de.bosch.com To transform the information contained in the context data into a latent representation of the target function , current approaches typically employ a mean aggregator and feed the output of this aggregator into a NN to predict a distribution over global latent parameters of the function . Hence , aggregation and latent parameter inference have so far been treated as separate parts of the learning pipeline . Moreover , when using a mean aggregator , every context sample is assumed to carry the same amount of information . Yet , in practice , different input locations have different task ambiguity and , therefore , samples should be assigned different importance in the aggregation process . In contrast , our Bayesian aggregation mechanism treats context aggregation and latent parameter inference as one holistic mechanism , i.e. , the aggregation directly yields the distribution over the latent parameters of the target function . Indeed , we formulate context aggregation as Bayesian inference of latent parameters using Gaussian conditioning in the latent space . Compared to existing methods , the resulting aggregator improves the handling of task ambiguity , as it can assign different variance levels to the context samples . This mechanism improves predictive performance , while it remains conceptually simple and introduces only negligible computational overhead . Moreover , our Bayesian aggregator can also be applied to deterministic model variants like the Conditional NP ( CNP ) ( Garnelo et al. , 2018a ) . In summary , our contributions are ( i ) a novel Bayesian Aggregation ( BA ) mechanism for context aggregation in NP-based models for probabilistic regression , ( ii ) its application to existing CLV architectures as well as to deterministic variants like the CNP , and ( iii ) an exhaustive experimental evaluation , demonstrating BA ’ s superiority over traditional mean aggregation . 2 RELATED WORK . Prominent approaches to probabilistic regression are Bayesian linear regression and its kernelized counterpart , the Gaussian process ( GP ) ( Rasmussen and Williams , 2005 ) . The formal correspondence of GPs with infinite-width Bayesian NNs ( BNNs ) has been established in Neal ( 1996 ) and Williams ( 1996 ) . A broad range of research aims to overcome the cubic scaling behaviour of GPs with the number of context points , e.g. , through sparse GP approximations ( Smola and Bartlett , 2001 ; Lawrence et al. , 2002 ; Snelson and Ghahramani , 2005 ; Quiñonero-Candela and Rasmussen , 2005 ) , by deep kernel learning ( Wilson et al. , 2016 ) , by approximating the posterior distribution of BNNs ( MacKay , 1992 ; Hinton and van Camp , 1993 ; Gal and Ghahramani , 2016 ; Louizos and Welling , 2017 ) , or , by adaptive Bayesian linear regression , i.e. , by performing inference over the last layer of a NN which introduces sparsity through linear combinations of finitely many learned basis functions ( Lazaro-Gredilla and Figueiras-Vidal , 2010 ; Hinton and Salakhutdinov , 2008 ; Snoek et al. , 2012 ; Calandra et al. , 2016 ) . An in a sense complementary approach aims to increase the data-efficiency of deep architectures by a fully Bayesian treatment of hierarchical latent variable models ( “ DeepGPs ” ) ( Damianou and Lawrence , 2013 ) . A parallel line of research studies probabilistic regression in the multi-task setting . Here , the goal is to formulate models which are data-efficient on an unseen target task by training them on data from a set of related source tasks . Bardenet et al . ( 2013 ) ; Yogatama and Mann ( 2014 ) , and Golovin et al . ( 2017 ) study multi-task formulations of GP-based models . More general approaches of this kind employ the meta-learning framework ( Schmidhuber , 1987 ; Thrun and Pratt , 1998 ; Vilalta and Drissi , 2005 ) , where a model ’ s training procedure is formulated in a way which incentivizes it to learn how to solve unseen tasks rapidly with only a few context examples ( “ learning to learn ” , “ few-shot learning ” ( Fei-Fei et al. , 2006 ; Lake et al. , 2011 ) ) . A range of such methods trains a meta-learner to learn how to adjust the parameters of the learner ’ s model ( Bengio et al. , 1991 ; Schmidhuber , 1992 ) , an approach which has recently been applied to few-shot image classification ( Ravi and Larochelle , 2017 ) , or to learning data-efficient optimization algorithms ( Hochreiter et al. , 2001 ; Li and Malik , 2016 ; Andrychowicz et al. , 2016 ; Chen et al. , 2017 ; Perrone et al. , 2018 ; Volpp et al. , 2019 ) . Other branches of meta-learning research aim to learn similarity metrics to determine the relevance of context samples for the target task ( Koch et al. , 2015 ; Vinyals et al. , 2016 ; Snell et al. , 2017 ; Sung et al. , 2017 ) , or explore the application of memory-augmented neural networks for meta-learning ( Santoro et al. , 2016 ) . Finn et al . ( 2017 ) propose model-agnostic meta-learning ( MAML ) , a general framework for fast parameter adaptation in gradient-based learning methods . A successful formulation of probabilistic regression as a few-shot learning problem in a multi-task setting is enabled by recent advances in the area of probabilistic meta-learning methods which allow a quantitative treatment of the uncertainty arising due to task ambiguity , a feature particularly relevant for few-shot learning problems . One line of work specifically studies probabilistic extensions of MAML ( Grant et al. , 2018 ; Ravi and Larochelle , 2017 ; Rusu et al. , 2018 ; Finn et al. , 2018 ; Kim et al. , 2018 ) . Further important approaches are based on amortized inference in multi-task CLV models ( Heskes , 2000 ; Bakker and Heskes , 2003 ; Kingma and Welling , 2013 ; Rezende et al. , 2014 ; Sohn et al. , 2015 ) , which forms the basis of the Neural Statistician proposed by Edwards and Storkey ( 2017 ) and of the NP model family ( Garnelo et al. , 2018b ; Kim et al. , 2019 ; Louizos et al. , 2019 ) . Gordon et al . ( 2019 ) present a unifying view on many of the aforementioned probabilistic architectures . Building on the conditional NPs ( CNPs ) proposed by Garnelo et al . ( 2018a ) , a range of NP-based architectures , such as Garnelo et al . ( 2018b ) and Kim et al . ( 2019 ) , consider combinations of deterministic and CLV model architectures . Recently , Gordon et al . ( 2020 ) extended CNPs to include translation equivariance in the input space , yielding state-of-the-art predictive performance . In this paper , we also employ a formulation of probabilistic regression in terms of a multi-task CLV model . However , while in previous work the context aggregation mechanism ( Zaheer et al. , 2017 ; Wagstaff et al. , 2019 ) was merely viewed as a necessity to consume context sets of variable size , we take inspiration from Becker et al . ( 2019 ) and emphasize the fundamental connection of latent parameter inference with context aggregation and , hence , base our model on a novel Bayesian aggregation mechanism . 3 PRELIMINARIES . We present the standard multi-task CLV model which forms the basis for our discussion and present traditional mean context aggregation ( MA ) and the variational inference ( VI ) likelihood approximation as employed by the NP model family ( Garnelo et al. , 2018a ; Kim et al. , 2019 ) , as well as an alternative Monte Carlo ( MC ) -based approximation . Problem Statement . We frame probabilistic regression as a multi-task learning problem . Let F denote a family of functions f ` : Rdx → Rdy with some form of shared statistical structure . We assume to have available data sets D ` ≡ { ( x ` , i , y ` , i ) } i of evaluations y ` , i ≡ f ` ( x ` , i ) + ε from a subset of functions ( “ tasks ” ) { f ` } L ` =1 ⊂ F with additive Gaussian noise ε ∼ N ( 0 , σ2n ) . From this data , we aim to learn the posterior predictive distribution p ( y ` |x ` , Dc ` ) over a ( set of ) y ` , given the corresponding ( set of ) inputs x ` as well as a context set Dc ` ⊂ D ` . The Multi-Task CLV Model . We formalize the multitask learning problem in terms of a CLV model ( Heskes , 2000 ; Gordon et al. , 2019 ) as shown in Fig . 1 . The model employs task-specific global latent variables z ` ∈ Rdz , as well as a task-independent latent variable θ , capturing the statistical structure shared between tasks . To learn θ , we split the data into context sets Dc ` ≡ { ( xc ` , n , yc ` , n ) } N ` n=1 and target sets Dt ` ≡ { ( xt ` , m , yt ` , m ) } M ` m=1 and maximize the posterior predictive likelihood function L∏ ` =1 p ( yt ` ,1 : M ` ∣∣xt ` ,1 : M ` , Dc ` , θ ) = L∏ ` =1 ∫ p ( z ` | Dc ` , θ ) M∏̀ m=1 p ( yt ` , m ∣∣ z ` , xt ` , m , θ ) dz ` ( 1 ) w.r.t . θ . In what follows , we omit task indices ` to avoid clutter . Likelihood Approximation . Marginalizing over the task-specific latent variables z is intractable for reasonably complex models , so one has to employ some form of approximation . The NP-family of models ( Garnelo et al. , 2018b ; Kim et al. , 2019 ) uses an approximation of the form log p ( yt1 : M ∣∣xt1 : M , Dc , θ ) ' Eqφ ( z|Dc∪Dt ) [ M∑ m=1 log p ( ytm ∣∣ z , xtm , θ ) + log qφ ( z| Dc ) qφ ( z| Dc ∪ Dt ) ] . ( 2 ) Being derived using a variational approach , this approximation utilizes an approximate posterior distribution qφ ( z| Dc ) ≈ p ( z| Dc , θ ) . Note , however , that it does not constitute a proper evidence lower bound for the posterior predictive likelihood since the intractable latent posterior p ( z| Dc , θ ) has been replaced by qφ ( z| Dc ) in the nominator of the rightmost term ( Le et al. , 2018 ) . An alternative approximation , employed for instance in Gordon et al . ( 2019 ) , also replaces the intractable latent posterior distribution by an approximate distribution qφ ( z| Dc ) ≈ p ( z| Dc , θ ) and uses a Monte-Carlo ( MC ) approximation of the resulting integral based on K latent samples , i.e. , log p ( yt1 : M ∣∣xt1 : M , Dc , θ ) ≈ − logK + log K∑ k=1 M∏ m=1 p ( ytm ∣∣ zk , xtm , θ ) , zk ∼ qφ ( z| Dc ) . ( 3 ) Note that both approaches employ approximations qφ ( z| Dc ) of the latent posterior distribution p ( z| Dc , θ ) and , as indicated by the notation , amortize inference in the sense that one single set of parameters φ is shared between all context data points . This enables efficient inference at test time , as no per-data-point optimization loops are required . As is standard in the literature ( Garnelo et al. , 2018b ; Kim et al. , 2019 ) , we represent qφ ( z| Dc ) and p ( ytm|z , xtm , θ ) by NNs and refer to them as the encoder ( enc , parameters φ ) and decoder ( dec , parameters θ ) networks , respectively . These networks set the means and variances of factorized Gaussian distributions , i.e. , qφ ( z| Dc ) = N ( z|µz , diag ( σ2z ) ) , µz = encµz , φ ( Dc ) , σ2z = encσ2z , φ ( D c ) , ( 4 ) p ( ytm ∣∣ z , xtm , θ ) = N ( ytm∣∣µy , diag ( σ2y ) ) , µy = decµy , θ ( z , xtm ) , σ2y = decσ2y , θ ( z , xtm ) . ( 5 ) Context Aggregation . The latent variable z is global in the sense that it depends on the whole context set Dc . Therefore , some form of aggregation mechanism is required to enable the encoder to consume context sets Dc of variable size . To represent a meaningful operation on sets , such an aggregation mechanism has to be invariant to permutations of the context data points . Zaheer et al . ( 2017 ) characterize possible aggregation mechanisms w.r.t . this permutation invariance condition , resulting in the structure of traditional aggregation mechanisms depicted in Fig . 2 ( a ) . Each context data tuple ( xcn , y c n ) is first mapped onto a latent observation rn = encr , φ ( x c n , y c n ) ∈ Rdr . Then , a permutation-invariant operation is applied to the set { rn } Nn=1 to obtain an aggregated latent observation r̄ . One prominent choice , employed for instance in Garnelo et al . ( 2018a ) , Kim et al . ( 2019 ) , and Gordon et al . ( 2019 ) , is to take the mean , i.e. , r̄ = 1 N N∑ n=1 rn . ( 6 ) Subsequently , r̄ is mapped onto the parameters µz and σ2z of the approximate posterior distribution qφ ( z| Dc ) using additional encoder networks , i.e. , µz = encµz , φ ( r̄ ) and σ2z = encσ2z , φ ( r̄ ) . Note that three encoder networks are employed here : ( i ) encr , φ to map from the context pairs to rn , ( ii ) encµz , φ to compute µz from the aggregated mean r̄ and ( iii ) encσ2z , φ to compute the variance σ 2 z from r̄ . In what follows , we refer to this aggregation mechanism as mean aggregation ( MA ) and to the networks encµz , φ and encσ2z , φ collectively as “ r̄-to-z-networks ” .

Summary: The paper builds upon previous lines of research on multi-task learning problem, such as conditional latent variable models including the Neural Process. As shown by the extensive Related Work section, this seems to be an active research direction. This makes it difficult for me to judge originality and significance, but it is well-written and clear.