Knowledge Flow: Improve Upon Your Teachers

A zoo of deep nets is available these days for almost any given task , and it is increasingly unclear which net to start with when addressing a new task , or which net to use as an initialization for fine-tuning a new model . To address this issue , in this paper , we develop knowledge flow which moves ‘ knowledge ’ from multiple deep nets , referred to as teachers , to a new deep net model , called the student . The structure of the teachers and the student can differ arbitrarily and they can be trained on entirely different tasks with different output spaces too . Upon training with knowledge flow the student is independent of the teachers . We demonstrate our approach on a variety of supervised and reinforcement learning tasks , outperforming fine-tuning and other ‘ knowledge exchange ’ methods . 1 INTRODUCTION . Research communities have amassed a sizable number of deep net architectures for different tasks , and new ones are added almost daily . Some of those architectures are trained from scratch while others are fine-tuned , i.e. , before training , their weights are initialized using a structurally similar deep net which was trained on different data . Beyond fine-tuning , particularly in reinforcement learning , teachers have also been considered in one way or another by Rusu et al . ( 2016b ) ; Fernando et al . ( 2017 ) ; Wang et al . ( 2017 ) ; Li & Hoiem ( 2016 ) ; Bengio et al . ( 2009 ) ; Patel et al . ( 2015 ) ; Chen & Liu ( 2016 ) ; Teh et al . ( 2017 ) ; Parisotto et al . ( 2016 ) . For instance , progressive neural net ( Rusu et al. , 2016b ) keeps multiple teachers during both training and inference , and learns to extract useful features from the teachers for a new target task . PathNet ( Fernando et al. , 2017 ) uses genetic algorithms to choose pathways from a giant network for learning new tasks . ‘ Growing a Brain ’ ( Wang et al. , 2017 ) fine-tunes a neural network while growing the network ’ s capacity ( wider or deeper layers ) . Actor-mimic ( Parisotto et al. , 2016 ) pre-trains a big model on multiple source tasks , then the big model is used as a weight initialization for a new model which will be trained on a new target task . Knowledge distillation ( Hinton et al. , 2015 ) distills knowledge from a large ensemble of models to a smaller student model . However , all the aforementioned techniques have limitations . For example , progressive neural net models ( Rusu et al. , 2016b ) grow with the number of teachers . This large number of parameters limits the number of teachers a progressive neural net can handle , and largely increases the training and testing time . In PathNet ( Fernando et al. , 2017 ) , searching over a big network for pathways is computationally intensive . For fine-tuning based methods such as ‘ Growing a Brain ’ ( Wang et al. , 2017 ) and actor-mimic ( Parisotto et al. , 2016 ) , only one pretrained model can be used at a time . Hence , their performance heavily relies on the chosen pretrained model . To address these shortcomings , we develop knowledge flow which moves ‘ knowledge ’ of multiple teachers when training a student . Irrespective of how many teachers we use , the student is guaranteed to become independent at the final stage of training and the size of the resulting student net remains constant . In addition , our framework makes no restrictions on the deep net size of the teacher and student , which provides flexibility in choosing teacher models . Importantly , our approach is applicable to a variety of tasks from reinforcement learning to fully-supervised training . We evaluate knowledge flow on a variety of tasks from reinforcement learning to fully-supervised learning . In particular , we follow Rusu et al . ( 2016b ) ; Fernando et al . ( 2017 ) and compare on the same Atari games . In addition , we also observed significant top-1 error rate improvements on supervised learning datasets , i.e. , CIFAR-10 , and CIFAR-100 . 2 BACKGROUND . Knowledge flow is applicable to a variety of settings from supervised learning to reinforcement learning , which we briefly review to introduce notation . Supervised Learning recovers the parameters θ of a mapping fθ : X → Y from data space X to output space Y . To this end , a dataset D = { ( xi , yi ) } ni=1 containing n pairs ( xi , yi ) ( assumed to be sampled i.i.d . ) is used , where xi ∈ X and yi ∈ Y . Given this dataset , the parameters θ of the mapping fθ are learned by minimizing a loss function ` ( x , y ) ( θ ) composed of a regularization term R ( θ ) and an empirical risk ` ( y , fθ ( x ) ) which compares groundtruth label y and prediction fθ ( x ) . The parameters θ are obtained by optimizing the following program : min θ E ( x , y ) ∼D [ ` ( x , y ) ( θ ) ] : = E ( x , y ) ∼D [ ` ( y , fθ ( x ) ) ] +R ( θ ) . ( 1 ) Hereby , the mapping fθ is obtained by maximizing the logits or a corresponding probability distribution f̂θ ( y|x ) , i.e. , fθ = argmaxy∈Y f̂θ ( y|x ) . Here and below let the hat ( ‘ ̂· ’ ) indicate probability distributions over appropriate domains . Reinforcement Learning considers an agent interacting with an environment according to a policy πθπ : X → A which maps a state xt ∈ X to an action at ∈ A at time t. The policy depends on the parameters θπ . After performing action at , the agent observes the next state xt+1 and receives a scalar reward rt . The discounted return at time t is defined as Rt = ∑∞ k=0 γ krt+k , where γ is the discount factor . The expected future reward when observing state x and when following policy πθπ is defined as V πθπ ( xt ) = Eτ∼πθπ [ Rt|xt ] , where τ = { ( xt , at , rt ) , ( xt+1 , at+1 , rt+1 ) , . . . } is a trajectory generated by following πθπ from state xt . The goal of reinforcement learning is to find a policy that maximizes the expected future reward from each state xt . Without loss of generality , in this paper , we follow the asynchronous advantage actor-critic ( A3C ) formulation ( Mnih et al. , 2016 ) . In A3C , the policy mapping πθπ ( x ) = argmaxa∈A π̂θπ ( a|x ) is obtained from a probability distribution over states , where π̂θπ ( a|x ) is modeled by a deep net with parameters θπ . The value function is also approximated by a deep net Vθv ( x ) , having parameters θv . To optimize the policy parameters θπ given a state xt , a loss function based on a scaled negative log-likelihood and a negative entropy regularizer is common : ` τπ ( θπ ) = 1 |τ | ∑ t∈τ [ − log π̂θπ ( at|xt ) ( Rt − Vθv ( xt ) ) − βH ( π̂θπ ( ·|xt ) ) ] . Hereby , Rt = ∑k−1 i=0 γ irt+i + γ kVθv ( xt+k ) is the empirical k-step return obtained when starting in state xt , and |τ | is the length of the trajectory τ generated by following πθπ . The scalar β ≥ 0 is a user-specified constant , and H ( π̂θπ ( ·|xt ) ) is the entropy function , which encourages exploration by favoring a uniform probability distribution π̂θπ ( a|x ) . To optimize the value function Vθv , it is common to use the squared loss ` τv ( θv ) = 1 2|τ | ∑ t∈τ ( Rt − Vθv ( xt ) ) 2 . By minimizing the empirical expectation of ` τπ ( θπ ) and ` τ v ( θv ) , i.e. , by addressing min θπ Eτ∼πθπ [ ` τ π ( θπ ) ] , and min θv Eτ∼πθπ [ ` τ v ( θv ) ] , ( 2 ) alternatingly , we learn a policy and a value function that maximize expected return . 3 KNOWLEDGE FLOW . Instead of optimizing the programs given in Eq . ( 1 ) and Eq . ( 2 ) from scratch , the aforementioned warm-start techniques ( see Sec . 5 for more ) are applicable . To address their mentioned shortcomings , we propose knowledge flow , a framework that moves ‘ knowledge ’ from an arbitrary number of deep nets , henceforth referred to as ‘ teachers ’ to a deep net under training , called the ‘ student . ’ 3.1 OVERVIEW . Knowledge flow is outlined on example deep nets in Fig . 1 ( a , b ) . We train the parameters of the student net which are randomly initialized . To this end we take advantage of teachers , whose parameters are fixed and obtained from pre-trained models on different source tasks by different algorithms . For example , for reinforcement learning , we may consider teachers trained by A3C ( Mnih et al. , 2016 ) , A2C ( Dhariwal et al. , 2017 ) or DQN ( Mnih et al. , 2015 ) . ‘ Knowledge ’ of multiple teachers is transferred to a student by adding transformed and scaled intermediate representations from the teacher deep nets to the student net . To achieve this , we modify the student net , i.e. , fθ in the supervised setting and πθπ ( a|x ) , Vθv ( x ) in the reinforcement learning case . We add teacher representations which are transformed by multiplication with a trainable matrix Q and scaled via a weight pw that is normalized to sum to one for each student layer and parameterized via trainable parameters w. The normalized weights encode which of the teachers ’ or the student ’ s representation to trust at every layer of the student net . Note that a teacher can help the student at different levels of abstraction with input from different levels of its net . Importantly , after training , the student model should perform well on the target task without relying on teachers . To achieve this , as training progresses , we increasingly encourage a high normalized weight on the student representation , which forces the student to eventually capture all the ‘ knowledge. ’ Due to the trainable scaling , at an early stage of training , we observe the student to rely heavily on the ‘ knowledge ’ of the teacher to quickly obtain better performance . However , as training proceeds , the student is encouraged to become more and more independent . During final stages of training , the student will no longer be able to rely on teachers , which ensures that the student has learned to master the desired task on its own . This is observed in Fig . 1 ( c ) . To formally encourage this successive transfer we introduce two additional loss functions . The first , referred to as the dependency loss ` dep ( w ) , captures how much a student relies on teachers . It depends on the weight vector w which encodes the strength of the coupling . The second one ensures that a student ’ s behavior doesn ’ t change rapidly when the teachers ’ influence decreases . We use loss ` KL ( · , · ) to capture the change . By combining student net modifications and additional loss terms , for the supervised task we obtain min θ , w , Q E ( x , y ) [ ˜̀ ( x , y ) ( θ , w , Q ) + λ1 ` dep ( w ) + λ2 ` KL ( ˜̂ fθ , ˜̂ fθold ) ] , ( 3 ) and for reinforcement learning the transformed program reads as follows : { minθπ , w , Q Eτ∼π̃θπ [ ˜̀ τ π ( θπ , w , Q ) + λ1 ` dep ( w ) + λ2 ` τ KL ( ˜̂πθπ , ˜̂πθπold ) ] minθv , w , Q Eτ∼π̃θπ [ ˜̀ τ v ( θv , w , Q ) ] . ( 4 ) Loss ˜̀·· ( θ , w , Q ) originates from the original loss ` · · ( θ ) ( Eqs . ( 1 ) - ( 2 ) ) by transforming the deep net to include cross-connections , hence its dependence on w , Q . The tilde ( ‘ ̃· ’ ) denotes this dependence , also for probability distribution ˜̂f and policy distribution ˜̂π . Parameters from the current and a previous iteration are referred to via θ and θold respectively . For both supervised and reinforcement learning , λ1 and λ2 control the strength which is used to decrease the influence of the teacher . A low λ1 allows the student to rely on teachers . Close to the end of training , the student should be independent . Therefore , we set λ1 to a small value at the beginning , and gradually increase its value as training progresses . Note that we don ’ t make any assumptions about teachers and student ’ s objective . If a teacher ’ s and student ’ s objective differ , negative transfer may occur initially . However , the proposed method quickly decreases the weight for teacher layers to reduce this effect . Despite differences , students could potentially still benefit from the low level representation of the teachers . We do observe this low level knowledge transfer in our experiments . In the following we first describe how to modify the deep nets , before we detail the loss functions ` dep and ` KL , which are used to successively decrease the influence of the teachers .

Summary: This paper presents a method for distilling multiple teacher networks into a student, by linearly combining feature representations from all networks at multiple intermediate layers, and gradually forcing the student to "take over" the learned combination.  Networks to be used as teachers are first pretrained on various initial tasks.  A student network is then trained on a target task (possibly different from any teacher task), by combining corresponding hidden layers from each teacher using learned linear remappings and weighted combinations.  Learning this combination allows the system to find appropriate teachers for the target task; eventually, a penalty on the combination weights forces all weight onto the student network, resulting in the distillation.