Transferring Inductive Biases through Knowledge Distillation

1 INTRODUCTION . Inductive biases are the characteristics of learning algorithms that influence their generalization behavior , independent of data . They are one of the main driving forces to push learning algorithms toward particular solutions ( Mitchell , 1980 ) . Having the right inductive biases is especially important for obtaining high performance when data or compute is a limiting factor , or when training data is not perfectly representative of the conditions at test time . Moreover , in the absence of strong inductive biases , a model can be equally attracted to several local minima on the loss surface ; and the converged solution can be arbitrarily affected by random variations like the initial state or the order of training examples ( Sutskever et al. , 2013 ; McCoy et al. , 2020 ; Dodge et al. , 2020 ) . There are different ways to inject inductive biases into learning algorithms , for instance through architectural choices , the objective function , the curriculum , or the optimization regime . In this paper , we exploit the power of Knowledge Distillation ( KD ) to transfer the effect of inductive biases between neural networks . KD refers to the process of transferring knowledge from a teacher model to a student model , where the logits from the teacher are used to train the student . KD is best known as an effective method for model compression ( Buciluǎ et al. , 2006 ; Hinton et al. , 2015 ; Sanh et al. , 2019 ) which allows taking advantage of a huge number of parameters during training while having an efficient smaller model during inference . The advantage of KD goes beyond model compression and it can be used to combine strengths of different learning algorithms ( Kuncoro et al. , 2019 ; 2020 ) . Different algorithms vary in terms of the computational/memory efficiency at training/inference or the inductive biases for learning particular patterns . This makes them better at solving certain problems and worse at others , i.e. , there is no “ one size fits all ” learning algorithm . Hence , it is important to explore the potential of KD for finding better trade-offs . The question that we ask in this paper is : “ In KD , are the preferences of the teacher that are rooted in its inductive biases , also reflected in its dark knowledge1 , and can they thus be transferred to the student ? ” . We are interested in cases where the student model can realize functions that are realizable by the teacher , i.e. , the student model is efficient with respect to the teacher model ( Cohen et al. , 2016 ) , while the teacher has a preference inductive bias so that the desired solutions are easily learnable for the teacher ( Seung et al. , 1991 ) . We consider two scenarios where the teacher and the student are neural networks with heterogeneous architectures , hence have different inductive biases . We train the models , both independently and using KD , on tasks for which having the right inductive biases is crucial . In the first test case , we study RNNs vs. Transformers ( Vaswani et al. , 2017 ) , on the subject-verb agreement prediction task ( Linzen et al. , 2016 ) . In this task , we use LSTMs ( Hochreiter & Schmidhuber , 1997 ) as the most widely used RNN variant . LSTMs are shown to perform better than vanilla Transformers in this task and their superior performance is attributed to their so-called “ recurrent ” inductive bias ( Tran et al. , 2018 ) . First , we identify the sources of the recurrent inductive bias of LSTMs : sequentiality , memory bottleneck , and recursion , and design experiments to show the benefits of each . Then , we show that through distilling knowledge of LSTMs to Transformers , the solutions that the Transformer models learn become more similar to the solution learned by LSTMs . In the second test case , we study CNNs vs. MLPs , in the context of the MNIST-C ( Corrupted MNIST ) benchmark ( Mu & Gilmer , 2019 ) , which is designed to measure out-of-distribution robustness of models . We train our models on MNIST and evaluate them on the Translated/Scaled MNIST . The particular form of parameter sharing in CNNs combined with the pooling mechanism makes them equivariant to these kinds of transformations ( Goodfellow et al. , 2016 ) , which leads to better generalization in these scenarios compared to MLPs . In our experiments and analysis on these two test cases2 , we compare the behavior of different models , from a wide range of perspectives , when trained in different setups including ( 1 ) when trained without KD , but directly from the data , ( 2 ) when trained with KD using a teacher with a similar architecture to the student , i.e . self-distillation , and ( 3 ) when trained with KD using a teacher with a different architecture that has stronger inductive biases that suit the task , compared to the student . As the first step , in setup ( 1 ) , i.e. , no KD , we demonstrate how inductive biases arising from different architectural choices affect the generalization behavior of the models we study ( §2.1 and §3.1 ) . We show that the models with more suitable inductive biases not only have better accuracy but also the solutions they converge to is a better solution in terms of other metrics . We also show that different instances of the model with stronger inductive biases have less variance in terms of all the metrics . Then , we apply KD to train the models and contrast the behavior of models trained with the setups ( 2 ) and ( 3 ) with the models trained with setup ( 1 ) , i.e . with KD vs. without KD . We show that regardless of the properties of the teacher , KD is a powerful technique in which the teacher model drives the student toward a particular set of solutions that is more restricted compared to the set of possible solutions that a student can converge to when it learns directly from data ( §2.2 , §3.2 , and Appendix C ) . 1Dark knowledge refers to the information encoded in the output logits of a neural network ( Hinton et al. , 2015 ) . 2The codes for the input pipelines , models , analysis , and the details of the hyper-parameters used in our experiments is available at https : //ANONYMIZED . Next , as the main contribution of our paper over previous works that study KD , we contrast the behavior of models trained with setup ( 3 ) with the models trained with setups ( 1 ) and ( 2 ) : • We show the performance of the student models in setup ( 3 ) increases , not only on in- distribution test sets ( §2.2 ) , but also on out-of-distribution data ( §3.2 ) . We demonstrate that this happens when the teacher has the right inductive bias and not necessarily otherwise , i.e. , setup ( 2 ) . • In setup ( 3 ) , besides performance , we show that , the solution that a student model converges to shares similar characteristics with the solution of its teacher . For instance in terms of confidence calibration ( §2.2 and §3.2 ) , and per sample behaviour of the model ( Appendix E ) . • We demonstrate that although the student model is merely exposed to the final logits of the teacher , the structure of the latent space of the student model becomes similar to the teacher , i.e . relational similarity of the internal representations from the student and its teacher increases ( §2.2 and §3.2 ) . As an example , in our second test case ( MNIST-C ) , when training an MLP model with KD using a CNN teacher , the student model explores the solution space in ways more similar to its teacher . Figure 1 visualizes and compares the path that an MLP takes during training ( Figure 1a ) , compared to a CNN ( Figure 1b ) . The CNN model explores the surface in a completely different manner than the MLP , while the path of a student MLP distilled from the CNN model as the teacher ( Figure1c ) is more similar to the CNN . 2 DISTILLING LSTMS INTO TRANSFORMERS . LSTMs and Transformers are the basic building blocks of many state-of-the-art models for sequence modeling and natural language processing . Transformers are an expressive class of models that do extremely well on many tasks where the training data is adequate in quantity ( Devlin et al. , 2019 ; Keskar et al. , 2019 ; Radford et al. , 2019 ; Brown et al. , 2020 ) . Several studies , however , have shown that LSTMs can perform better than Transformers on tasks requiring sensitivity to ( linguistic ) structure , especially when the data is limited ( Tran et al. , 2018 ; Dehghani et al. , 2019 ) . We chose the subject-verb agreement prediction task , introduced by Linzen et al . ( 2016 ) , as the test case , as it yields a meaningful difference between LSTMs and Transformers ( Tran et al. , 2018 ) . We compare these two families of models and conduct experiments to emphasize the differences between them when trained independently and through KD . Recurrent Inductive Bias . Among sequence modeling architectures , models with recursion are in particular powerful for natural language processing due to their adequacy to model hierarchical structures ( Linzen et al. , 2016 ) . The recursion in a model can be implemented in various ways , like in Recurrent Neural Networks ( Elman , 1990 ) , Recursive Neural Networks ( Socher et al. , 2010 ; Le & Zuidema , 2014 ) and Universal Transformers ( Dehghani et al. , 2019 ; Hao et al. , 2019 ) . While theoretically , both recurrent neural networks ( RNNs ) and Transformers can deal with finite hierarchical structures , empirical results indicate the superiority of RNNs over Transformers ( Tran et al. , 2018 ; Dehghani et al. , 2019 ; Hahn , 2020 ) . In the literature ( Sutskever et al. , 2013 ; Dehghani et al. , 2019 ) , the inductive bias of RNNs is referred to as the recurrent inductive bias . Here , we distinguish between three main sources of this bias : 1 . Sequentiality : There is an inherent notion of order in the architecture that forces the model to access the next tokens in the input one by one and process them sequentially . 2 . Memory bottleneck : The model has no direct access to the past tokens and has to compress all the information from the past in a hidden state , which is accessible when processing a new token . 3 . Recursion : The model recursively applies the same function on the varying input at every step . Transformers ( Vaswani et al. , 2017 ) , in contrast , process the input in parallel . Although a weak notion of order is encoded by positional embeddings , no explicit assumption is made in the connectivity structure of the architecture . Moreover , they have a global receptive field and can access all tokens through self-attention . Finally , standard Transformers are not recursive . However , the standard Transformer can be modified to have an architecture with specifications that are similar to RNNs . We provide empirical results to demonstrate the benefits of these different sources of inductive biases of RNNs . For this purpose , we design experiments with variants of Transformers in which we attempt to approximate some of the RNNs ’ assumptions . Task and Models . We study the performance of LSTMs and variants of Transformers on the task of predicting number-agreement between subjects and verbs in English sentences . We investigate the quality of the solutions they converge to when they are trained independently and when they are trained through distillation . We use the subject-verb agreement dataset of Linzen et al . ( 2016 ) , for which the size of the training set is ∼121k examples and the size of the test set is ∼1m . Succeeding at this task is a strong indicator that a model can learn syntactic structures and is therefore proposed by Linzen et al . ( 2016 ) as a proxy for assessing the ability of models to capture hierarchical structure in natural language . It is shown that RNNs have better inductive biases to learn this compared to standard Transformers ( Tran et al. , 2018 ; Dehghani et al. , 2019 ) . In this task , examples are grouped into different levels of difficulty based on the number of “ agreement attractors ” 3 , and distance between the verb and its subject . Hence , we report both micro accuracy ( µ−Accuracy ) and macro accuracy over different groups in terms of distance ( D−Accuracy ) and numbers of attractors ( A−Accuracy ) . Similar to Tran et al . ( 2018 ) , we follow two setups : 1 ) when the learning objective is next word prediction , i.e. , language modeling ( LM ) ; 2 ) when we directly optimize for predicting the verb number , singular or plural , i.e. , classification . In the LM setup , we look at the probabilities predicted when the target of the prediction is the verb of interest , and see whether the probability of the correct form of the verb is higher than the other form ( singular vs plural ) . In the classification setup , the input to the model is a sentence up to the position of the verb of interest and the model predicts whether the verb at that position is singular or plural . In the LM setup , we employ two unidirectional LSTMs with different sizes , LSTM and Small LSTM , and two Transformers , Transformer and Small Transformer . In this setup , corresponding LSTMs and Transformers have roughly the same number of parameters . In the classification setup we compare the following models : ( 1 ) a standard unidirectional LSTM ( sequentiality + memory bottleneck + recursion ) ( 2 ) Transformer : Transformer encoder with a class token ( CLS ) for classification , BERT ( Devlin et al. , 2019 ) style , ( 3 ) Transformer-seq : Transformer encoder with future masking where the classification is done using the representation of the last token4 ( sequentiality ) , ( 4 ) UniversalTransformer-seq : Universal Transformer ( Dehghani et al. , 2019 ) encoder , in which the parameters are shared in depth , with future masking ( sequentiality + recursion ) . Appendix H provides more details on the architectures .

Summary: The paper investigates the oft-overlooked aspect of knowledge distillation (KD) -- why it works. The paper highlights the ability of KD for transferring not just the soft labels, but the inductive bias (assumptions inherent in the method, e.g. LSTM's notion of sequentiality, and CNN's translational invariance/equivariance) from the student so that the student exhibits, to an extent, the teacher's generalization properties as well. The paper explores doing KD between LSTMs and several versions of Transformers (with varying structural constraints) on a subject-verb-agreement dataset, and between CNNs and MLPs on MNIST and corrupted MNIST. Compared to prior work showing that better teacher performance lead to better student performance, this paper also shows that the student's performance on different aspects becomes more similar to the teacher's -- (1) if the teacher is strong on metric A and weak on metric B compared to a student on its own, the student can become stronger on A and weaker on B when distilled using the teacher; (2) if the teacher can generalize well to a separate, previously unseen dataset but the student generalizes poorly on its own, after distillation the student can generalize much better than it can possibly learn to on its own.