Teaching with Commentaries

1 INTRODUCTION . Training , regularising , and understanding complex neural network models is challenging . There remain central open questions on making training faster and more data-efficient ( Kornblith et al. , 2019 ; Raghu et al. , 2019a ; b ) , ensuring better generalisation ( Zhang et al. , 2016 ) and improving transparency and robustness ( Bau et al. , 2017 ; Madry et al. , 2017 ) . A promising approach for addressing these questions is learning to teach ( Zhu , 2015 ) , in which learned auxiliary information about a task is provided to a neural network to inform the training process and help downstream objectives . Examples include providing auxiliary training targets ( Liu et al. , 2019 ; Navon et al. , 2020 ; Pham et al. , 2020 ) and reweighting training examples to emphasise important datapoints ( Fan et al. , 2020 ; Jiang et al. , 2018 ; Ren et al. , 2018 ; Shu et al. , 2019 ) . Learning to teach approaches have achieved promising results in vision and language applications ( Jiang et al. , 2018 ; Ren et al. , 2018 ; Shu et al. , 2019 ; Hu et al. , 2019 ) using a handful of specific modifications to the training process . In this paper , we take steps towards generalising these approaches , introducing a flexible and effective learning to teach framework using commentaries . Commentaries represent learned meta-information helpful for training a model on a task , and once learned , such commentaries can be reused as is to improve the training of new models . We demonstrate that commentaries can be used for applications ranging from speeding up training to gaining insights into the neural network model . Specifically , our contributions are : 1 . We formalise the notion of commentaries , providing a unified framework for learning metainformation that can be used to improve network training and examine model learning . 2 . We present gradient-based methods to learn commentaries by optimising a network ’ s validation loss , leveraging recent work in implicit differentiation to scale to larger models . 3 . We use commentaries to define example-weighting curricula , a common method of teaching neural networks . We show that these learned commentaries hold interpretable insights , lead to speedups in training , and improve performance on few-shot learning tasks . ∗Work done while interning at Google . 4 . We define data augmentation policies with label-dependent commentaries , and obtain insights into the design of effective augmentation strategies and improved performance on benchmark tasks as compared to baselines . 5 . We parameterise commentaries as attention masks to find important regions of images . Through qualitative and quantitative evaluation , we show these masks identify salient image regions and can be used to improve the robustness of neural networks to spurious background correlations . 6 . We show that learned commentaries can generalise : when training new models , reusing learned commentaries can lead to learning speed/performance improvements . This suggests a use-case for commentaries : being stored with a dataset and leveraged to improve training of new models . 2 TEACHING WITH COMMENTARIES . Definition : We define a commentary to be learned information helpful for ( i ) training a model on a task or ( ii ) providing insights on the learning process . We envision that commentaries , once learned , could be stored alongside a dataset and reused as is to assist in the training of new models . Appendix A explores a simple instantiation of commentaries for Celeb-A ( Liu et al. , 2015 ) , to provide intuition of the structures that commentaries can encode . Formally , let t ( x , y , i ; φ ) denote a commentary that is a function of a data point x , prediction target y , and iteration of training i , with parameters φ . The commentary may be represented in a tabular fashion for every combination of input arguments , or using a neural network that takes these arguments as inputs . The commentary is used to train a student network n ( x ; θ ) with parameters θ . 2.1 LEARNING COMMENTARIES . We now describe algorithms to learn commentaries 1 . Throughout , we denote the training set as DT , the validation set as DV and the loss function ( e.g . cross-entropy ) as L. With θ denoting the parameters of the student network and φ denoting the commentary parameters , we let θ̂ , φ̂ be the respective optimised parameters . We seek to find φ̂ such that the student network ’ s validation loss , LV , is minimised . As the commentary is used during the training of the student network , LV implicitly depends on φ , enabling the use of gradient-based optimisation algorithms to find φ̂ . Algorithm 1 : Backpropagation Through Training : When student network training has a small memory footprint , we optimise commentary parameters by iterating the following process , detailed in Algorithm 1 : ( 1 ) train a student and store the computation graph during training ; ( 2 ) compute the student ’ s validation loss ; ( 3 ) calculate the gradient of this loss w.r.t . the commentary parameters by backpropagating through training ; ( 4 ) update commentary parameters using gradient descent . By optimizing the commentary parameters over the entire trajectory of student learning , we encourage this commentary to be effective when used in the training of new student networks . This supports the goal of the commentary being stored with the dataset and reused in future model learning . Algorithm 2 : Large-Scale Commentary Learning with Implicit Differentiation : When training the student model has a large memory footprint , backpropagating through training to obtain exact commentary parameter gradients is too memory expensive . We therefore leverage the Implicit Function Theorem ( IFT ) and efficient inverse Hessian approximation to obtain approximate gradients , following Lorraine et al . ( 2020 ) . The gradient of the validation loss w.r.t . the commentary parameters can be expressed as : ∂LV ∂φ = ∂LV ∂θ̂ × ∂θ̂ ∂φ . ( 3 ) The first term on the right hand side in equation 3 is simple to compute , but the second term is expensive . Under fixed-point and regularity assumptions on student and commentary parameters ( θ̂ ( φ ) , φ ) , the IFT allows expressing this second term ∂θ̂∂φ as the following product : ∂θ̂ ∂φ = − [ ∂2LT ∂θ ∂θT ] −1 × ∂ 2LT ∂θ ∂φT ∣∣∣ θ̂ ( φ ) , ( 4 ) 1Code at https : //github.com/googleinterns/commentaries Algorithm 1 Commentary Learning through Backpropagation Through Training . 1 : Initialise commentary parameters φ 2 : for t = 1 , . . . , T meta-training steps do 3 : Initialise student network n ( x ; θ ) with parameters θ0 4 : Train student network with N steps of gradient descent to optimise : LT ( θ , φ ) = Ex , y∼DT [ L̃ ( n ( x ; θ ) , t ( · ; φ ) , y ) ] , ( 1 ) where L̃ is a loss function adjusted from L to incorporate the commentary , and LT ( θ , φ ) is the expected adjusted loss over the training data . Output : θ̂ , the optimised parameters of student network ( implicitly a function of φ , θ̂ ( φ ) ) . 5 : Compute validation loss : LV ( φ ) = Ex , y∼DV [ L ( n ( x ; θ̂ ( φ ) ) , y ) ] ( 2 ) 6 : Compute ∂LV ( φ ) ∂φ , by backpropagating through the N steps of student training , and update φ . 7 : end for 8 : Output : φ̂ , the optimised parameters of the commentary . Algorithm 2 Commentary Learning through Implicit Differentiation . 1 : Initialise commentary parameters φ and student network parameters θ 2 : for t = 1 , . . . , M do 3 : Compute the student network ’ s training loss , LT ( θ , φ ) , equation 1 . 4 : Compute the gradient of this loss w.r.t the student parameters θ . 5 : Perform a single gradient descent update on the parameters to obtain θ̂ ( implicitly a function of φ , θ̂ ( φ ) ) . 6 : Compute the student network ’ s validation loss , LV ( φ ) , equation 2 . 7 : Compute ∂LV ∂θ̂ . 8 : Approximately compute ∂θ̂∂φ with equation 4 , using a truncated Neumann series with a single term and implicit vector-Jacobian products ( Lorraine et al. , 2020 ) . 9 : Compute the overall derivative ∂LV∂φ using steps ( 7 ) and ( 8 ) , and update φ . 10 : Set θ ← θ̂ . 11 : end for 12 : Output : φ̂ , the optimised parameters of the commentary . i.e. , a product of an inverse Hessian and a matrix of mixed partial derivatives . Following Lorraine et al . ( 2020 ) , we efficiently approximate this product using a truncated Neumann series and implicit vector-Jacobian products . Leveraging this approximation then yields a second method for commentary learning , described in Algorithm 2 . Since a single term in the Neumann series is sufficient for learning , each iteration of this algorithm has similar time complexity to a single iteration of training . In this method , commentary parameters are learned jointly with student parameters , avoiding training a single student model multiple times . This approach therefore scales to millions of commentary parameters and large student models ( Hataya et al. , 2020 ; Lorraine et al. , 2020 ) . However , since the commentary is not directly optimised over the entire trajectory of learning , its generalisability to new models is not ensured . We examine this in our experiments , demonstrating that commentaries learned in this manner can indeed generalise to training new student networks . 3 COMMENTARIES FOR EXAMPLE WEIGHTING CURRICULA . We now explore our first main application of commentaries : encoding a separate weight for each training example at each training iteration . Since the commentaries are a function of the training iteration , they can encode curriculum structure , so we refer to them as curriculum commentaries . We specify these weights using a commentary neural network ( or teacher network ) t ( x , i ; φ ) → [ 0 , 1 ] that produces a weight for every training example at every iteration of training of the student network . When training a student network , using the notation of §2.1 , the commentary is incorporated in the training loss as : L̃ = t ( x , i ; φ ) · L ( n ( x ; θ ) , y ) , where L ( · ) is the original loss function for the task . The validation loss is unweighted . 3.1 SYNTHETIC EXAMPLE : ROTATED MNIST DIGITS . We first learn example weight curriculum commentaries on a synthetic MNIST binary classification problem . Each example in the dataset is a rotated MNIST digit ‘ 1 ’ , with variable rotation angle that defines the class . We generate two datasets : the non-overlapping dataset and the overlapping dataset . In the non-overlapping dataset , the rotation angle for each example from class 1 and class 0 is drawn from non-overlapping distributions Uniform [ 15 , 45 ] and Uniform [ −45 , −15 ] respectively . In the overlapping dataset , the rotation angles are drawn from overlapping distributions Uniform [ −5 , 30 ] and Uniform [ −30 , 5 ] respectively ( Figure 1 ) . We use two block CNNs as both the commentary neural network and student network . The commentary network takes as input the image and the iteration of student training , and outputs a weight for each example in the batch . We learn commentary parameters by backpropagating through student training ( Algorithm 1 , §2.1 ) , and use 500 gradient steps for inner optimisation ( i.e. , N = 500 ) . Implementation is with the higher library ( Grefenstette et al. , 2019 ) . Further details in Appendix B.1 . Results : Figure 1 visualises the two datasets and plots the learned example weights as a function of rotation at iteration 500 of the student training . When classes do not overlap ( left ) , the example weights are highest for those examples near to the decision boundary ( small rotation magnitude ) . When the classes do overlap ( right ) , the more representative examples further from the boundary are upweighted and ambiguous examples in the overlap region are downweighted : a sensible result . We perform further analysis of the learned example weighting curriculum in Appendix B.1 , demonstrating that the learned curricula in both cases are meaningful . Overall , these results demonstrate that the learned commentaries capture interesting and intuitive structure .

Summary: This paper proposes a general framework for boosting CNNs performance on different tasks by using'commentary' to learn meta-information. The obtained meta-information can also be used for other purposes such as the mask of objects within spurious background and the similarities among classes. The commentary module would be incorporated into standard networks and be iteratively optimized with the host via the proposed objective. To effectively optimize both the commentary and the standard network, this paper adopts the techniques including implicit function theorem and efficient inverse Hessian approximations.