Meta Learning via Learned Loss

1 INTRODUCTION Inspired by the remarkable capability of humans to quickly learn and adapt to new tasks , the concept of learning to learn , or meta-learning , recently became popular within the machine learning community ( Andrychowicz et al. , 2016 ; Duan et al. , 2016 ; Finn et al. , 2017 ) . We can classify learning-to-learn methods into roughly 2 categories : approaches that lead to learning representations that can generalize and are easily adaptable to new tasks ( Finn et al. , 2017 ) , and learning approaches that attempt to learn how to optimize models ( Andrychowicz et al. , 2016 ; Duan et al. , 2016 ) . In this paper we investigate the second type of approach and propose a learning framework that is able to learn loss function representations that can then be used to optimize models for new tasks . Specifically , the purpose of this work is to encode learning strategies into an adaptive high-dimensional loss function , or a meta-loss , which generalizes across multiple training contexts or tasks . Inspired by inverse reinforcement learning ( Ng et al. , 2000 ) , our work combines the learning to learn paradigm of meta-learning with the generality of learning loss landscapes . We construct a unified , fully differentiable framework that can learn model-agnostic loss functions to provide a strong learning signal for a large range of model classes , such as classifiers , regressors or control policies . The contributions of this work are as follows : i ) we present a framework for learning adaptive , high-dimensional loss functions through back-propagation that shape the loss landscape such that it can be efficiently optimized with gradient descent . This framework involves an inner and an outer loop . In the inner loop , a model or an optimizee is trained with gradient descent using the loss coming from our learned meta-loss function . Fig . 1 shows the pipeline for updating the optimizee with the meta-loss . The outer loop optimizes the meta-loss function by minimizing a task-loss , such as a standard regression or reinforcement-learning loss , that is induced by the updated optimizee . We show that our learned meta-loss functions improves over directly learning via the task-loss itself while maintaining the generality of the task-loss . ii ) We show how we can utilize extra information that helps shape the loss landscapes at meta-train time . This extra information can take on various forms , such as exploratory signals or expert demonstrations for RL tasks . After training the meta-loss function , the task-specific losses are no longer required since the training of optimizees can be performed entirely by using the meta-loss function alone , without requiring the extra information given at meta-train time . In this way , our meta-loss can find more efficient ways to optimize the original task loss . 2 RELATED WORK . Meta-learning originates in the concept of learning to learn ( Schmidhuber , 1987 ; Bengio & Bengio , 1990 ; Thrun & Pratt , 2012 ) . Let us consider gradient based learning approaches , that update the parameters of an optimizee fθ ( x ) , with model parameters θ and inputs x as follows : θnew = hψ ( θ , ∇θLφ ( y , fθ ( x ) ) ; ( 1 ) where we take the gradient of a loss function L , parametrized by φ , with respect to the optimizee ’ s parameters θ and use a gradient transform h , parametrized by ψ , to compute new model parameters θnew 1 . In this context , we can divide related work on meta-learning into learning model parameters θ that can be easily adapted to new tasks ( Finn et al. , 2017 ; Mendonca et al. , 2019 ; Gupta et al. , 2018 ; Yu et al. , 2018 ) , learning optimizer policies h that transform parameters updates with respect to known loss or reward functions ( Maclaurin et al. , 2015 ; Andrychowicz et al. , 2016 ; Li & Malik , 2016 ; Franceschi et al. , 2017 ; Meier et al. , 2018 ; Duan et al. , 2016 ) , or learning loss/reward function representations φ ( Sung et al. , 2017 ; Houthooft et al. , 2018 ; Zou et al. , 2019 ) . Alternatively , in unsupervised learning settings , meta-learning has been used to learn unsupervised rules that can be transferred between tasks ( Metz et al. , 2019 ; Hsu et al. , 2018 ) Our framework falls into the category of learning loss landscapes . Similar to works by Sung et al . ( 2017 ) and Houthooft et al . ( 2018 ) , we aim at learning loss function parameters φ that can be applied to various optimizee models , e.g . regressors , classifiers or agent policies . Our learned loss functions are independent of the model parameters θ that are to be optimized , thus they can be easily transferred to other optimizee models . This is in contrast to methods that meta-learn model-parameters θ directly , such as ( Finn et al. , 2017 ; Mendonca et al. , 2019 ) , where the learned representation θ can not be separated from the original model of the optimizee . The idea of learning loss landscapes or reward functions in the reinforcement learning ( RL ) setting can be traced back to the field of inverse reinforcement learning ( Ng et al. , 2000 ; Abbeel & Ng , 2004 , IRL ) . However , in contrast to IRL we do not require expert demonstrations ( however we can incorporate them ) . Instead we use task losses as a measure of the effectiveness of our loss function when using it to optimize an optimizee , and as the optimization objective for φ ( parameters of the learned loss ) . Closest to our method are the works on evolved policy gradients ( Houthooft et al. , 2018 ) , teacher networks ( Wu et al. , 2018 ) , meta-critics ( Sung et al. , 2017 ) and meta-gradient RL ( Xu et al. , 2018 ) . In contrast to using an evolutionary approach ( e.g . Houthooft et al. , 2018 ) , we design a differentiable framework and describe a way to optimize the loss function with gradient descent in both supervised and reinforcement learning settings . Wu et al . ( 2018 ) propose that instead of learning a differentiable loss function directly , a teacher network is trained to predict parameters of a manually designed loss function , which requires a new teacher network design and training for each new loss function class . In Xu et al . ( 2018 ) , discount and bootstrapping parameters are learned online to optimize a task-specific meta-objective . Our method does not require manual design of the loss function parameterization or choosing particular parameters that have to be optimized , as our loss functions are learned entirely from data . Finally , in work by Sung et al . ( 2017 ) a meta-critic is learned to provide a task-conditional value function , used to train an actor policy . Although training a meta-critic in the supervised setting reduces to learning a loss function as in our work , in the reinforcement learning setting we show that it is possible to use learned loss functions to optimize policies directly with gradient descent . 3 META-LEARNING VIA LEARNED LOSS . In this work , we aim to learn a loss function , which we call meta-loss , that is used to train an optimizee , e.g . a classifier , a regressor or an agent policy . More concretely , we aim to learn a meta-loss function Mφ with parameters φ , which can be used to predict the loss value Llearned that can be used to train an optimizee fθ with parameters θ via gradient descent : θnew = θ − α∇θLlearned , where Llearned =Mφ ( y , fθ ( x ) ) ( 2 ) where y can be ground target information in supervised learning settings or goal and state information for reinforcement learning settings . In short , we aim to learn a loss function that can be used as 1for simple gradient descent h ( θ , ∇θL ( y , fθ ( x ) ) = θ − ψ∇θL ( y , fθ ( x ) ) Algorithm 1 ML3 at ( meta-train ) 1 : φ , θ ← randomly initialize 2 : while not done do 3 : x , y ← Sample task samples from T 4 : Llearned =M ( y , fθ ( x ) ) 5 : θnew ← θ − α∇θ Ex [ Llearned ] 6 : φ← φ− η∇φLT ( y , fθnew ) Algorithm 2 ML3 at ( meta-test ) 1 : M ← # of optimizee updates 2 : θ ← randomly initialize 3 : for j ∈ { 0 , . . . , M } do 4 : x , y ← Sample task samples from T 5 : Llearned =M ( y , fθ ( x ) ) 6 : θ ← θ − α∇θ Ex [ Llearned ] depicted in Algorithm 2 . Towards this goal , we propose an algorithm to train learning loss function parameters φ via gradient descent . During meta-train time the goal is to train the loss function parameters φ such that it successfully optimizes a model fθ for chosen task ( s ) T . The key challenge is to derive a training signal for loss parameters φ . In the following , we describe our approach to addressing this challenge , which we call Meta-Learning via Learned Loss ( ML3 ) . 3.1 ML3 FOR SUPERVISED LEARNING We start with supervised learning settings , in which our framework aims at learning a meta-loss functionMφ ( y , fθ ( x ) ) that predicts the loss value given the ground truth target y and the predicted target fθ ( x ) . For clarity purposes we constrain the following presentation to learning a meta-loss network that predicts the loss value for training a regressor fθ via gradient descent , however the methodology trivially generalizes to classification tasks . Our meta-learning framework starts with randomly initialized model parameters θ and loss parameters φ . The current loss parameters are then used to predict loss value Llearned = Mφ ( y , fθ ( x ) ) . To optimize model parameters θ we need to compute the gradient of the loss value with respect to θ , ∇θL = ∇θMφ ( y , fθ ( x ) ) . Using the chain rule , we can decompose the gradient computation into the gradient of the loss network with respect to predictions of model fθ ( x ) times the gradient of model f with respect to model parameters2 , ∇θMφ ( y , fθ ( x ) ) = ∇fMφ ( y , fθ ( x ) ) ∇θfθ ( x ) . ( 3 ) Once we have updated the model parameters θnew = θ − α∇θLlearned using the current meta-loss network parameters φ , we want to measure how much learning progress has been made with lossparameters φ and optimize φ via gradient descent . Note , that the new model parameters θnew are implicitly a function of loss-parameters φ , because changing φ would lead to different θnew . In order to evaluate θnew , and through that loss-parameters φ , we introduce the notion of a task-loss during meta-train time . For instance , we use the Mean-Squared-Error loss , which is typically used for regression tasks , as a task-loss LT = ( y − fθnew ( x ) ) 2 . We now optimize loss parameters φ by taking the gradient of LT with respect to φ as follows2 : ∇φLT ( y , fθnew ( x ) ) = ∇θnewLT∇φθnew ( 4 ) = ∇θnewLT∇φ [ θ − α∇θE [ Mφ ( x , fθ ( x ) ) ] ( 5 ) Optimization of the loss-parameters can either happen after each inner gradient step ( where inner refers to using the current loss parameters to update θ ) , or after M inner gradient steps with the current meta-loss networkMφ . The latter option requires back-propagation through a chain of all optimizee update steps . In practice we notice that updating the meta-parameters φ after each inner gradient update step works better . We summarize the meta-train phase in Algorithm 1 , with one inner gradient step . 3.2 ML3 REINFORCEMENT LEARNING In this section , we introduce several modifications that allow us to apply the ML3 framework to reinforcement learning problems . Let M = ( S , A , P , R , p0 , γ , T ) be a finite-horizon Markov Decision Process ( MDP ) , where S and A are state and action spaces , P : S × A × S → R+ is a state-transition probability function or system dynamics , R : S × A → R a reward function , p0 : S → R+ an initial state distribution , γ a reward discount factor , and T a horizon . Let 2alternatively this gradient computation can performed using automatic differentiation τ = ( s0 , a0 , . . . , sT , aT ) be a trajectory of states and actions and R ( τ ) = ∑T−1 t=0 γ tR ( st , at ) the trajectory return . The goal of reinforcement learning is to find parameters θ of a policy πθ ( a|s ) that maximizes the expected discounted reward over trajectories induced by the policy : Eπθ [ R ( τ ) ] where s0 ∼ p0 , st+1 ∼ P ( st+1|st , at ) and at ∼ πθ ( at|st ) . In what follows , we show how to train a meta-loss network to perform effective policy updates in a reinforcement learning scenario . To apply our ML3 framework , we replace the optimizee fθ from the previous section with a stochastic policy πθ ( a|s ) . We present two applications of ML3 to RL . 3.2.1 ML3 FOR MODEL-BASED REINFORCEMENT LEARNING Model-based RL ( MBRL ) attempts to learn a policy π by first learning a dynamic model P . Intuitively , if the model P is accurate , we can use it to optimize the policy parameters θ . As we typically do not know the dynamics model a-priori , MBRL algorithms iterate between using the current approximate dynamics model P , to optimize the policy π such that it maximizes the reward R under P , then use the optimized policy π to collect more data which is used to update the model P . In this context , we aim to learn a loss function that is used to optimize policy parameters through our meta-networkM . Similar to the supervised learning setting we use current meta-parameters φ to optimize policy parameters θ under the current dynamics model P : θnew = θ − α∇θ [ Mφ ( τ , g ) ] , where τ = ( s0 , a0 , . . . , sT , aT ) is the sampled trajectory and the variable g captures some task-specific information , such as the goal state of the agent . To optimize φ we again need to define a task loss , which in the MBRL setting can be defined as LT ( g , πθnew ) = −Eπθnew , P [ Rg ( τnew ) ] , denoting the reward that is achieved under the current dynamics model P . To update φ , we compute the gradient of the task loss LT wrt . φ , which involves differentiating all the way through the reward function , dynamics model and the policy that was updated using the meta-lossMφ . The pseudo-code in Algorithm 3 ( Appendix A.1 ) illustrates the MBRL learning loop . In Algorithm 5 ( Appendix A.1 ) , we show the policy optimization procedure during meta-test time . Notably , we have found that in practice , the model of the dynamics P is not needed anymore for policy optimization at meta-test time . The meta-network learns to implicitly represent the gradients of the dynamics model and can produce a loss to optimize the policy directly . 3.2.2 ML3 FOR MODEL-FREE REINFORCEMENT LEARNING Finally , we consider the model-free reinforcement learning ( MFRL ) case , where we learn a policy without learning a dynamics model . In this case , we can define a surrogate objective , which is independent of the dynamics model , as our task-specific loss ( Williams , 1992 ; Sutton et al. , 2000 ; Schulman et al. , 2015 ) : LT ( g , πθnew ) = −Eπθnew [ Rg ( τnew ) log πθnew ( τnew ) ] = −Eπθnew [ Rg ( τnew ) T−1∑ t=0 log πθnew ( at|st ) ] ( 6 ) Similar to the MBRL case , the task loss is indirectly a function of the meta-parameters φ that are used to update the policy parameters . Although we are evaluating the task loss on full trajectory rewards , we perform policy updates from Eq . 2 using stochastic gradient descent ( SGD ) on the meta-loss with mini-batches of experience ( si , ai , ri ) for i ∈ { 0 , . . . , B−1 } with batch size B , similar to Houthooft et al . ( 2018 ) . The inputs of the meta-loss network are the sampled states , sampled actions , task information g and policy probabilities of the sampled actions : Mφ ( s , a , πθ ( a|s ) , g ) . In this way , we enable efficient optimization of very high-dimensional policies with SGD provided only with trajectory-based rewards . In contrast to the above MBRL setting , the rollouts used for task-loss evaluation are real system rollouts , instead of simulated rollouts . At test time , we use the same policy update procedure as in the MBRL setting , see Algorithm 5 ( Appendix A.1 ) . 3.3 THE TASK LOSS AND ADDING EXTRA INFORMATION DURING meta-train So far , we have discussed using standard task losses , such as MSE-loss for regression or reward functions for RL settings . However , it is possible to provide more information about the task at metatrain time , which can influence the learning of the loss-landscape . In our work , we experiment with 3 different types of extra information at meta-train time : for supervised learning we show that providing ground truth information about the optimal parameters can help shape a convex loss-landscape for otherwise non-convex optimization problems ; for reinforcement learning tasks we demonstrate that by providing additional rewards in the task loss during meta-train time , we can encourage the trained meta-loss to learn exploratory behaviors ; and finally also for reinforcement learning tasks we also show how expert demonstrations can be incorporated to learn loss functions which can generalize to new tasks . In all these settings , the additional information shapes the learned loss function such that the environment does not need to provide this information during meta-test time .

Summary: This paper proposes a new meta-learning method (ML3) that meta-learns a loss function that is able to generalize across tasks. Building upon bi-level optimization framework as in MAML, instead of using a task-specific loss function in the inner loop, the authors compute adapted parameters of the model using a parametrized loss network and learn the loss network via backpropagation. Experiments are conducted on supervised sinusoid regression and binary digit classification as well as on model-based and model-free RL benchmarks.