One Network Fits All? Modular versus Monolithic Task Formulations in Neural Networks

1 INTRODUCTION . Standard practice in machine learning has long been to only address carefully circumscribed , often very related tasks . For example , we might train a single classifier to label an image as containing objects from a certain predefined set , or to label the words of a sentence with their semantic roles . Indeed , when working with relatively simple classes of functions like linear classifiers , it would be unreasonable to expect to train a classifier that handles more than such a carefully scoped task ( or related tasks in standard multitask learning ) . As techniques for learning with relatively rich classes such as neural networks have been developed , it is natural to ask whether or not such scoping of tasks is inherently necessary . Indeed , many recent works ( see Section 1.2 ) have proposed eschewing this careful scoping of tasks , and instead training a single , “ monolithic ” function spanning many tasks . Large , deep neural networks can , in principle , represent multiple classifiers in such a monolithic learned function ( Hornik , 1991 ) , giving rise to the field of multitask learning . This combined function might be learned by combining all of the training data for all of the tasks into one large batch–see Section 1.2 for some examples . Taken to an extreme , we could consider seeking to learn a universal circuit—that is , a circuit that interprets arbitrary programs in a programming language which can encode various tasks . But , the ability to represent such a monolithic combined function does not necessarily entail that such a function can be efficiently learned by existing methods . Cryptographic hardness theorems ( Kearns & Valiant , 1994 ) establish that this is not possible in general by any method , let alone the specific training methods used in practice . Nevertheless , we still can ask how ∗Work performed in part while visiting Google . †Work performed in part while affiliated with Stanford , and in part while interning at Google . rich a family of tasks can be learned by these standard methods . In this work , we study the extent to which backpropagation with stochastic gradient descent ( SGD ) can learn such monolithic functions on diverse , unrelated tasks . There might still be some inherent benefit to an architecture in which tasks are partitioned into sub-tasks of such small scope , and the training data is correspondingly partitioned prior to learning . For example , in the early work on multitask learning , Caruana ( 1997 ) observed that training a network to solve unrelated tasks simultaneously seemed to harm the overall performance . Similarly , the seminal work of Jacobs et al . ( 1991 ) begins by stating that “ If backpropagation is used to train a single , multilayer network to perform different subtasks on different occasions , there will generally be strong interference effects that lead to slow learning and poor generalization ” . We therefore ask if , for an unfortunate choice of tasks in our model , learning by standard methods might be fundamentally impaired . As a point of reference from neuroscience , the classical view is that distinct tasks are handled in the brain by distinct patches of the cortex . While it is a subject of debate whether modularity exists for higher level tasks ( Samuels , 2006 ) , it is accepted that there are dedicated modules for low-level tasks such as vision and audio processing . Thus , it seems that the brain produces a modular architecture , in which different tasks are handled by different regions of the cortex . Conceivably , this division into task-specific regions might be driven by fundamental considerations of learnability : A single , monolithic neural circuit might simply be too difficult to learn because the different tasks might interfere with one another . Others have taken neural networks trained by backpropagation as a model of learning in the cortex ( Musslick et al. , 2017 ) ; to the extent that this is reasonable , our work has some bearing on these questions as well . 1.1 OUR RESULTS . We find , perhaps surprisingly , that combining multiple tasks into one can not fundamentally impair learning with standard training methods . We demonstrate this for a broad family of methods for combining individual tasks into a single monolithic task . For example , inputs for each individual tasks may come from a disjoint region ( for example , a disjoint ball ) in a common input space , and each individual task could then involve applying some arbitrary simple function ( e.g. , a separate linear classifier for each region ) . Alternately there may be an explicit “ task code ” attribute ( e.g. , a one-hot code ) , together with the usual input attributes and output label ( s ) , where examples with the same task code are examples for the same learning task . Complementing our results that combining multiple tasks does not impair learning , we also find that some task coding schemes do incur a sample complexity penalty . A vast variety of task coding schemes may be used . As a concrete example , when the data points for each task are well-separated into distinct clusters , and the tasks are linear classification tasks , we show that a two-layer architecture trained with SGD successfully learns the combined , monolithic function ; the required amount of data simply scales as the sum of the amount required to learn each task individually ( Theorem 2 ) . Meanwhile , if the tasks are determined by a balanced decision tree of height h on d code attributes ( as in Fig . 1 , left ) , we find that the training time and amount of data needed scales as ∼ dh—quasipolynomial in the 2h leaves ( distinct tasks ) when d is of similar size to h , and thus when the coding is efficient ( Theorem 3 ) . We also prove a corresponding lower bound , which shows that this bound is in fact asymptotically tight ( Theorem 3 ) . More generally , for task codings based on decision trees using linear splits with a margin of at least γ ( when the data has unit ` 2 norm ) , the training time and required data are asymptotically bounded by ∼ eO ( h/γ 2 ) , which for constant γ is polynomial in the 2h functions ( Theorem 4 ) . We generalize from these cluster-based and decision-tree based task codings to more complex codes that are actually simple programs . For instance , we show that SQL-style aggregation queries over a fixed database , written as a functions of the parameters of the query , can also be learned this way . More generally , simple programming constructs ( such as in Fig . 1 , right ) , built by operations such as compositions , aggregation , concatenation , and branching on a small number of such learnable functions , are also learnable ( Theorem 5 ) . In general , we can learn a low-depth formula ( circuit with fan-out 1 ) in which each gate is not merely a switch ( as in a decision tree ) , but can be any analytic function on the inputs , including arithmetic operations . Again , our key technical contribution is that we show that all of these functions are efficiently learned by SGD . This is non-trival since , although universal approximation theorems show that such functions can be expressed by ( sufficiently wide ) two-layer neural networks , under standard assumptions some expressible functions are not learnable Klivans & Sherstov ( 2009 ) . We supplement the theoretical bounds with experiments on clusters , decision trees , and SQL-style aggregation showing that such functions are indeed learned in practice . We note that the learning of such combined functions could have been engineered by hand : for example , there exist efficient algorithms for learning clusterings or such decision trees , and it is easy to learn the linear classifiers given the partitioned data . Likewise , these classes of functions are all known to be learnable by other methods , given an appropriate transformation of the input features . The key point is that the two-layer neural network can jointly learn the task coding scheme and the task-specific functions without special engineering of the architecture . That is , it is unnecessary to engineer a way of partitioning of the data into separate tasks prior to learning . Relatedly , the time and sample requirements of learning multiple tasks on a single network in general is insufficient to explain the modularity observed in biological neural networks if their learning dynamics are similar to SGD —i.e. , we can not explain the presence of modularity from such general considerations . All our theoretical results are based upon a fundamental theorem that shows that analytic functions can be efficiently learnt by wide ( but finite-width ) two-layer neural networks with standard activation functions ( such as ReLU ) , using SGD from a random initialization . Specifically , we derive novel generalization bounds for multivariate analytic functions ( Theorems 1 and 8 ) by relating wide networks to kernel learning with a specific network-induced kernel ( Jacot et al. , 2018 ; Du et al. , 2019 ; Allen-Zhu et al. , 2019 ; Arora et al. , 2019a ; Lee et al. , 2019 ) , known as the neural tangent kernel ( NTK ) ( Jacot et al. , 2018 ) . We further develop a calculus of bounds showing that the sum , product , ratio , and composition of analytic functions is also learnable , with bounds constructed using the familiar product and chain rules of univariate calculus ( Corollaries 1 , 2 ) . These above learnability results may be of independent interest ; for example , they can be used to show that natural physical laws like the gravitational force equations ( shown in Fig . 1 ) can be efficiently learnt by neural networks ( Section B.1 ) . Furthermore , our bounds imply that the NTK kernel for ReLU activation has theoretical learning guarantees that are superior to the Gaussian kernel ( Section A.2 ) , which we also demonstrate empirically with experiments on learning the gravitational force law ( Section B.2 ) . 1.2 RELATED WORK . Most related to our work are a number of works in application areas that have sought to learn a single network that can perform many different tasks . In natural language processing , Tsai et al . ( 2019 ) show that a single model can solve machine translation across more than 50 languages . Many other works in NLP similarly seek to use one model for multiple languages , or even multiple tasks ( Johnson et al. , 2017 ; Aharoni et al. , 2019 ; Bapna et al. , 2019 ; Devlin et al. , 2018 ) . Monolithic models have also been successfully trained for tasks in very different domains , such as speech and language ( Kaiser et al. , 2017 ) . Finally , there is also work on training extremely large neural networks which have the capacity to learn multiple tasks ( Shazeer et al. , 2017 ; Raffel et al. , 2019 ) . These works provide empirical clues that suggest that a single network can successfully be trained to perform a wide variety of tasks . But , they do not provide a systematic theoretical investigation of the extent of this ability as we do here . Caruana ( 1997 ) proposed multitask learning in which a single network is trained to solve multiple tasks on the same input simultaneously , as a vector of outputs . He observed that average generalization error for the multiple tasks may be much better than when the tasks are trained separately , and this observation initiated an active area of machine learning research ( Zhang & Yang , 2017 ) . Multitask learning is obviously related to our monolithic architectures . The difference is that whereas in multitask learning all of the tasks are computed simultaneously and output on separate gates , here all of the tasks share a common set of outputs , and the task code inputs switch between the various tasks . Furthermore , contrary to the main focus of multitask learning , we are primarily interested in the extent to which different tasks may interfere , rather than how much similar ones may benefit . Our work is also related to studies of neural models of multitasking in cognitive science . In particular , Musslick et al . ( 2017 ) consider a similar two-layer architecture in which there is a set of task code attributes . But , as in multitask learning , they are interested in how many of these tasks can be performed simultaneously , on distinct outputs . They analyze the tradeoff between improved sample complexity and interference of the tasks with a handcrafted “ gating ” scheme , in which the parts of activity are zeroed out depending on the input ( as opposed to the usual nonlinearities ) ; in this model , they find out that the speedup from multitask learning comes at the penalty of limiting the number of tasks that can be correctly computed as the similarity of inputs varies . Thus , in contrast to our model where the single model is computing distinct tasks sequentially , they do find that the distinct tasks can interfere with each other when we seek to solve them simultaneously .

Summary: This paper takes an interesting theoretical dive on the learnability of multiple tasks as one task, while the tasks are constructed with special structures of the cluster, decision tree, or simple program. This paper provides sample complexities analysis, showing that wide two-layer neural networks with standard activation functions and SGD optimization is able to capture the data regularity between input and output, modulated by the three kinds of task code considered. Experiment results show that the networks are flexible enough to fit complex data generated in this way.