Original: MA-BERT: T OWARDS M ATRIX A RITHMETIC - ONLY BERT I NFERENCE BY E LIMINATING C OMPLEX N ONLINEAR F UNCTIONS
Revised: MA-BERT: T OWARDS M ATRIX A RITHMETIC - ONLY BERT I NFERENCE BY E LIMINATING C OMPLEX N ONLINEAR F UNCTIONS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Due to their superior results, Transformer-based models such as BERT have become de facto standards in many Natural Language Processing (NLP) applications.
Revised: Due to their superior results, Transformer-based models such as BERT have become de facto standards in many Natural Language Processing (NLP) applications.

Original: However, the intensive use of complex non-linear functions within the Transformer architecture impairs its computing efficiency and complicates corresponding accelerator designs, because non-linear functions are generally computation-intensive and require special hardware support.
Revised: However, the intensive use of complex non-linear functions within the Transformer architecture impairs its computing efficiency and complicates corresponding accelerator designs, because non-linear functions are generally computation-intensive and require special hardware support.

Original: In light of this, we propose MA-BERT, which allows matrix arithmetic-only operations in
Revised: In light of this, we propose MA-BERT, which allows matrix arithmetic-only operations in

Original: Transformer-based NLP models and achieves efficient inference with negligible accuracy loss.
Revised: Transformer-based NLP models and achieves efficient inference with negligible accuracy loss.

Original: Specifically, we propose four correlated techniques that include approximating softmax with a two-layer neural network, replacing GELU with
Revised: Specifically, we propose four correlated techniques that include approximating softmax with a two-layer neural network, replacing GELU with

Original: ReLU, fusing normalization layers with adjacent linear layers, and leveraging knowledge transfer from baseline models.
Revised: ReLU, fusing normalization layers with adjacent linear layers, and leveraging knowledge transfer from baseline models.

Original: Through these techniques, we are able to eliminate the major non-linear functions in Transformer-based models and obtain MA-BERT with only matrix arithmetic and trivial ReLU operations without compromising on accuracy.
Revised: Through these techniques, we are able to eliminate the major non-linear functions in Transformer-based models and obtain MA-BERT with only matrix arithmetic and trivial ReLU operations without compromising on accuracy.

Original: With mainly regular matrix arithmetic operations,
Revised: With mainly regular matrix arithmetic operations,

Original: MA-BERT enables hardware-friendly processing on various computing engines, including CPUs, GPUs, and customized neural network accelerators.
Revised: MA-BERT enables hardware-friendly processing on various computing engines, including CPUs, GPUs, and customized neural network accelerators.

Original: Our experimental results on CPUs show that MA-BERT achieves up to 27% reduction in inference time with comparable accuracy on many downstream tasks compared to the baseline BERT models.
Revised: Our experimental results on CPUs show that MA-BERT achieves up to 27% reduction in inference time with comparable accuracy on many downstream tasks compared to the baseline BERT models.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Recently, pretrained Transformer-based models such as GPT (Radford et al., 2018) and BERT (Devlin et al., 2018), have consistently dominated the leaderboards for a variety of NLP tasks and even surpassed the human baseline.
Revised: Recently, pretrained Transformer-based models such as GPT (Radford et al., 2018) and BERT (Devlin et al., 2018), have consistently dominated the leaderboards for a variety of NLP tasks and even surpassed the human baseline.

Original: Consequently, there has been a strong push for these models to be used in many NLP applications.
Revised: Consequently, there has been a strong push for these models to be used in many NLP applications.

Original: At the same time, they have become a popular research topic in academia and industry.
Revised: At the same time, they have become a popular research topic in academia and industry.

Original: This in turn led to the creation of even better models such as GPT3 (Brown et al., 2020), RoBERTa (Liu et al., 2019), and DeBERTa (He et al., 2020), which further stretched the limits of what such models can achieve.
Revised: This in turn led to the creation of even better models such as GPT3 (Brown et al., 2020), RoBERTa (Liu et al., 2019), and DeBERTa (He et al., 2020), which further stretched the limits of what such models can achieve.

Original: Clearly, the advent of Transformer-based models has brought about a paradigm shift in the NLP field from the pre-Transformer era when recurrent neural networks and their variants used to dominate.
Revised: Clearly, the advent of Transformer-based models has brought about a paradigm shift in the NLP field from the pre-Transformer era when recurrent neural networks and their variants used to dominate.

Original: Nonetheless, the exceptional performance of these Transformer-based models partly stems from their deep structure which involves a huge number of parameters and operations during a single forward propagation.
Revised: Nonetheless, the exceptional performance of these Transformer-based models partly stems from their deep structure which involves a huge number of parameters and operations during a single forward propagation.

Original: These characteristics make it challenging to meet timing requirements and implement on devices with limited memory and computational power.
Revised: These characteristics make it challenging to meet timing requirements and implement on devices with limited memory and computational power.

Original: Consequently, many notable works have been proposed to improve their inference performance.
Revised: Consequently, many notable works have been proposed to improve their inference performance.

Original: In particular to BERT, some works seek to compress the model by reducing the number of encoder layers via knowledge distillation (Sanh et al., 2019;
Revised: In particular to BERT, some works seek to compress the model by reducing the number of encoder layers via knowledge distillation (Sanh et al., 2019;

Original: Aguilar et al., 2020; Sun et al., 2020; 2019;
Revised: Aguilar et al., 2020; Sun et al., 2020; 2019;

Original: Jiao et al., 2019; Xu et al., 2020) with a minor accuracy penalty.
Revised: Jiao et al., 2019; Xu et al., 2020) with a minor accuracy penalty.

Original: Other works focus on model pruning (Gao et al., 2021; Gordon et al., 2020; Voita et al., 2019; Wang et al., 2021), which involves eliminating unessential weights or attention heads, leading to sparser models with a reduced number of parameters and computing operations.
Revised: Other works focus on model pruning (Gao et al., 2021; Gordon et al., 2020; Voita et al., 2019; Wang et al., 2021), which involves eliminating unessential weights or attention heads, leading to sparser models with a reduced number of parameters and computing operations.

Original: Additionally, several works leverage quantization (Kim et al., 2021; Bhandare et al.,
Revised: Additionally, several works leverage quantization (Kim et al., 2021; Bhandare et al.,

Original: ; Zafrir et al., 2019) to reduce the precision of weights and computing operations, which lowers memory demands and speeds up inference.
Revised: ; Zafrir et al., 2019) to reduce the precision of weights and computing operations, which lowers memory demands and speeds up inference.

Original: Despite these works, one non-negligible overhead in Transformer-based models is often overlooked – the intensive use of complex non-linear functions.
Revised: Despite these works, one non-negligible overhead in Transformer-based models is often overlooked – the intensive use of complex non-linear functions.

Original: In BERT, these non-linear functions include the softmax operation, Gaussian Error Linear Unit (GELU) activation function (Hendrycks & Gimpel, 2016), and Layer Normalization (LayerNorm) (Ba et al., 2016).
Revised: In BERT, these non-linear functions include the softmax operation, Gaussian Error Linear Unit (GELU) activation function (Hendrycks & Gimpel, 2016), and Layer Normalization (LayerNorm) (Ba et al., 2016).

Original: Although these functions undeniably play a role in helping the model to learn better during training time, they become a considerable bottleneck during inference as they are not straightforward to evaluate.
Revised: Although these functions undeniably play a role in helping the model to learn better during training time, they become a considerable bottleneck during inference as they are not straightforward to evaluate.

Original: Furthermore, in hardware accelerator designs, these non-linear functions typically require a separate hardware for acceleration (Liu et al., 2021; Khan et al., 2021), which makes it challenging to be deployed on resourceconstrained computing platforms.
Revised: Furthermore, in hardware accelerator designs, these non-linear functions typically require a separate hardware for acceleration (Liu et al., 2021; Khan et al., 2021), which makes it challenging to be deployed on resourceconstrained computing platforms.

Original: In Figure 1 and Table 1, we show the breakdown of cycle budget for the BERT base model.
Revised: In Figure 1 and Table 1, we show the breakdown of cycle budget for the BERT base model.

Original: For non-linear operations, the cycle budget is defined as the equivalent number of cycles that the matrix multiply takes to process, given the matrix multiply dimensions and the number of multiplications per cycle (Khan et al., 2021).
Revised: For non-linear operations, the cycle budget is defined as the equivalent number of cycles that the matrix multiply takes to process, given the matrix multiply dimensions and the number of multiplications per cycle (Khan et al., 2021).

Original: From the figure, we can see that those non-linear operations  take around 44% of the processing time.
Revised: From the figure, we can see that those non-linear operations can take as much as 43% of the processing time.

Original: The computation efficiency of BERT has a huge potential to improve if we can optimize these non-linear operations.
Revised: The computation efficiency of BERT has a huge potential to improve if we can optimize these non-linear operations.

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: In view of this, we introduce an efficient BERT, termed MA-BERT, which makes a novel attempt to completely eliminate the complex non-linear functions in BERT by substituting them with simpler functions.
Revised: In view of this, we introduce an efficient BERT, termed MA-BERT, which makes a novel attempt to completely eliminate the complex non-linear functions in BERT by substituting them with simpler functions.

Original: In MA-BERT, four correlated techniques are applied: (1) approximating softmax with a two-layer neural network, (2) replacing GELU with ReLU, (3) fusing normalization layers with adjacent linear layers, and (4) leveraging knowledge transfer from pretrained baseline models.
Revised: In MA-BERT, four correlated techniques are applied: (1) approximating softmax with a two-layer neural network, (2) replacing GELU with ReLU, (3) fusing normalization layers with adjacent linear layers, and (4) leveraging knowledge transfer from pretrained baseline models.

Original: Through these techniques, MA-BERT achieves matrix arithmetic-only and trivial ReLU operations with negligible accuracy loss.
Revised: Through these techniques, MA-BERT achieves matrix arithmetic-only and trivial ReLU operations with negligible accuracy loss.

Original: Our experiments show that the performance of MA-BERT on downstream tasks is on par with the corresponding BERT baseline and yet achieves up to 27% performance improvement on general
Revised: Our experiments show that the performance of MA-BERT on downstream tasks is on par with the corresponding BERT baseline and yet achieves up to 27% performance improvement on general

Original: CPUs.
Revised: CPUs.

Original: By eliminating the complex non-linear functions, the majority of the processing in MABERT becomes regular matrix-matrix arithmetic and can be potentially deployed on various computing engines including CPUs, GPUs, and neural network processing units (NPUs) without any hardware modification.
Revised: By eliminating the complex non-linear functions, the majority of the processing in MABERT becomes regular matrix-matrix arithmetic and can be potentially deployed on various computing engines including CPUs, GPUs, and neural network processing units (NPUs) without any hardware modification.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: While so many works focus on model compression techniques, only some paid attention to the optimization of non-linear functions in BERT and other variants of Transformer-based models.
Revised: While so many works focus on model compression techniques, only some paid attention to the optimization of non-linear functions in BERT and other variants of Transformer-based models.

Original: For instance, in MobileBERT, Sun et al.
Revised: For instance, in MobileBERT, Sun et al.

Original: (2020) proposed a thin version of BERT large with reduced hidden dimensions.
Revised: (2020) proposed a thin version of BERT large with reduced hidden dimensions.

Original: As part of their operational optimization, they replaced GELU with ReLU and removed LayerNorm to further reduce MobileBERT’s latency.
Revised: As part of their operational optimization, they replaced GELU with ReLU and removed LayerNorm to further reduce MobileBERT’s latency.

Original: However, their results showed that the inclusion of these optimizations degrades the model’s performance.
Revised: However, their results showed that the inclusion of these optimizations degrades the model’s performance.

Original: Kim et al.
Revised: Kim et al.

Original: (2021), on the other hand, quantized the non-linear functions in BERT to 32-bit integer precision, allowing these non-linear functions to be computed using faster integer arithmetic.
Revised: (2021), on the other hand, quantized the non-linear functions in BERT to 32-bit integer precision, allowing these non-linear functions to be computed using faster integer arithmetic.

Original: Specific to tackling the latency overhead of LayerNorm, Zhang & Sennrich (2019) proposed root mean square layer normalization which reduces the amount of computation compared to LayerNorm by only using root mean square statistics to normalize inputs.
Revised: Specific to tackling the latency overhead of LayerNorm, Zhang & Sennrich (2019) proposed root mean square layer normalization which reduces the amount of computation compared to LayerNorm by only using root mean square statistics to normalize inputs.

Original: In the field of computer vision, Yao et al.
Revised: In the field of computer vision, Yao et al.

Original: (2021) reintroduced Batch Normalization (BatchNorm) (Ioffe & Szegedy, 2015) as a replacement for LayerNorm (Ba et al., 2016) into vision Transformers and demonstrated that the resulting architecture performs on par with its LayerNorm-based counterpart.
Revised: (2021) reintroduced Batch Normalization (BatchNorm) (Ioffe & Szegedy, 2015) as a replacement for LayerNorm (Ba et al., 2016) into vision Transformers and demonstrated that the resulting architecture performs on par with its LayerNorm-based counterpart.

Original: Similarly, Shen et al.
Revised: Similarly, Shen et al.

Original: (2020) proposed PowerNorm as a replacement for LayerNorm in Transformers and demonstrated that PowerNorm can outperform LayerNorm in neural machine translation and language modelling tasks.
Revised: (2020) proposed PowerNorm as a replacement for LayerNorm in Transformers and demonstrated that PowerNorm can outperform LayerNorm in neural machine translation and language modelling tasks.

Original: Due to its proven performance and use of batch statistics, we have decided to adopt PowerNorm in our paper to minimize the overhead introduced by the normalization operations in BERT during inference.
Revised: Due to its proven performance and use of batch statistics, we have decided to adopt PowerNorm in our paper to minimize the overhead introduced by the normalization operations in BERT during inference.

Original: Leaning onto the hardware side, Yu et al.
Revised: Leaning onto the hardware side, Yu et al.

Original: (2022) proposed a general approximation framework which enabled them to convert all the non-linear operations into table look-ups.
Revised: (2022) proposed a general approximation framework which enabled them to convert all the non-linear operations into table look-ups.

Original: This helped to reduce the cost of non-linear functions but incur new type of operations and corresponding hardware.
Revised: This helped to reduce the cost of non-linear functions but incur new type of operations and corresponding hardware.

Original: Different from these works, our work strives to completely eliminate the complex non-linear functions in BERT and achieve matrix arithmetic-only operations with trivial ReLU, which could benefit inference on both general computing units and accelerator designs for edge applications.
Revised: Different from these works, our work strives to completely eliminate the complex non-linear functions in BERT and achieve matrix arithmetic-only operations with trivial ReLU, which could benefit inference on both general computing units and accelerator designs for edge applications.

Original:  M ETHOD
Revised:  M ETHOD

Original: This section gives a brief overview of the non-linear functions that exist in BERT and highlights their computation inefficiency during inference.
Revised: This section gives a brief overview of the non-linear functions that exist in BERT and highlights their computation inefficiency during inference.

Original: Following that, more efficient methods are proposed.
Revised: Following that, more efficient methods are proposed.

Original: 3. S OFTMAX : A PPROXIMATING WITH A 2-L AYER N EURAL N ETWORK
Revised: 3. S OFTMAX : A PPROXIMATING WITH A 2-L AYER N EURAL N ETWORK

Original: Every encoder layer in BERT contains a multi-head self-attention (MHA) sublayer followed by a feed forward network (FFN) sublayer.
Revised: Every encoder layer in BERT contains a multi-head self-attention (MHA) sublayer followed by a feed forward network (FFN) sublayer.

Original: The MHA sublayer endows the model the ability to capture various dependencies between different parts of an input sequence (Vaswani et al., 2017) and is defined as follows:MultiHead
Revised: The MHA sublayer endows the model the ability to capture various dependencies between different parts of an input sequence (Vaswani et al., 2017) and is defined as follows:

Original: [Equation]
Revised: [Equation]

Original: where headi = softmax ( XW
Revised: where head

Original: [Equation]
Revised: [Equation]

Original: In each self-attention head, a s × s scaled dot product of the query and key matrices is passed through a softmax function to obtain the weights on the value matrix.
Revised: In each self-attention head, a s × s scaled dot product of the query and key matrices is passed through a softmax function to obtain the weights on the value matrix.

Original: Here, s denotes the sequence length.
Revised: Here, s denotes the sequence length.

Original: The softmax function is computed as:
Revised: The softmax function is computed as:

Original: [Equation]
Revised: [Equation]

Original: From Eq.
Revised: From Eq.

Original: 3 it is clear that all inputs along the sequence dimension have to be processed before the softmax value of a particular input can be computed.
Revised: 3 it is clear that all inputs along the sequence dimension have to be processed before the softmax value of a particular input can be computed.

Original: This becomes prohibitively expensive as the input sequence length increases (Du et al., 2019; Stevens et al., 2021).
Revised: This becomes prohibitively expensive as the input sequence length increases (Du et al., 2019; Stevens et al., 2021).

Original: Moreover, the softmax function contains expensive exponentiation and division operations which further complicate accelerator designs (Hu et al., 2018).
Revised: Moreover, the softmax function contains expensive exponentiation and division operations which further complicate accelerator designs (Hu et al., 2018).

Original: In view of this, we decided to find an alternative for the softmax function.
Revised: In view of this, we decided to find an alternative for the softmax function.

Original: The universal approximation theorem (Hornik et al., 1989) states that a neural network with a single hidden layer and non-linear activation function can approximate any function, given sufficient hidden neurons.
Revised: The universal approximation theorem (Hornik et al., 1989) states that a neural network with a single hidden layer and non-linear activation function can approximate any function, given sufficient hidden neurons.

Original: As such, we propose to substitute the softmax function in each layer of BERT with a simple 2-layer neural network.
Revised: As such, we propose to substitute the softmax function in each layer of BERT with a simple 2-layer neural network.

Original: ReLU is chosen as the activation function of the hidden layer with hardware reuse in mind since we will also be adopting it in the FFN sublayer.
Revised: ReLU is chosen as the activation function of the hidden layer with hardware reuse in mind since we will also be adopting it in the FFN sublayer.

Original: Note that this modification will inevitably fix the maximum input sequence length of our modified model to the input layer size of the neural network.
Revised: Note that this modification will inevitably fix the maximum input sequence length of our modified model to the input layer size of the neural network.

Original: To eliminate the effects of any padding tokens in an input sequence, we multiplied the scaled dot product of the query and key matrices with the input attention mask before feeding it as an input to the neural network.
Revised: To eliminate the effects of any padding tokens in an input sequence, we multiplied the scaled dot product of the query and key matrices with the input attention mask before feeding it as an input to the neural network.

Original: Given that the same softmax function is being approximated, it should also be possible to approximate softmax using the same neural network across the encoder layers, allowing us to reduce the parameters introduced.
Revised: Given that the same softmax function is being approximated, it should also be possible to approximate softmax using the same neural network across the encoder layers, allowing us to reduce the parameters introduced.

Original: 3. GELU: R EPLACING WITH R E LU
Revised: 3. GELU: R EPLACING WITH R E LU

Original: Unlike the original Transformer architecture proposed by Vaswani et al.
Revised: Unlike the original Transformer architecture proposed by Vaswani et al.

Original: (2017), many Transformerbased NLP models such as BERT use the GELU activation function instead of ReLU in their FFN layers to enhance their learning capability.
Revised: (2017), many Transformerbased NLP models such as BERT use the GELU activation function instead of ReLU in their FFN layers to enhance their learning capability.

Original: GELU is defined as follows:
Revised: GELU is defined as follows:

Original: [Equation]
Revised: [Equation]

Original: According to Eq. 4, it is clear that the GELU involves complex operations such as exponentiation and integration which are computation-intensive and hard to be implemented efficiently in accelerator designs.
Revised: According to Eq. 4, it is clear that the GELU involves complex operations such as exponentiation and integration which are computation-intensive and hard to be implemented efficiently in accelerator designs.

Original: One possible way is to employ a look up table to simplify the operations, but doing so entails a trade-off between accuracy and hardware utilization.
Revised: One possible way is to employ a look up table to simplify the operations, but doing so entails a trade-off between accuracy and hardware utilization.

Original: Other implementations involving function approximations (Hendrycks & Gimpel, 2016; Kim et al., 2021) also exist, but they often involve multiple operations which ends up consuming a notable portion of inference time (Pati et al., 2021).
Revised: Other implementations involving function approximations (Hendrycks & Gimpel, 2016; Kim et al., 2021) also exist, but they often involve multiple operations which ends up consuming a notable portion of inference time (Pati et al., 2021).

Original: In contrast, ReLU simply outputs zero or the input value depending on the input’s sign, and can be implemented easily in hardware using a single multiplexer.
Revised: In contrast, ReLU simply outputs zero or the input value depending on the input’s sign, and can be implemented easily in hardware using a single multiplexer.

Original: With the potential simplification in hardware and speed-up during inference (Sun et al., 2020), we propose to bring ReLU back as a replacement for GELU.
Revised: With the potential simplification in hardware and speed-up during inference (Sun et al., 2020), we propose to bring ReLU back as a replacement for GELU.

Original: 3. L AYER N ORM : F USING N ORMALIZATION L AYERS WITH A DJACENT L INEAR L AYERS
Revised: 3. L AYER N ORM : F USING N ORMALIZATION L AYERS WITH A DJACENT L INEAR L AYERS

Original: Every encoder layer in BERT involves two LayerNorm sublayers – one after the MHA sublayer and another after the FFN sublayer.
Revised: Every encoder layer in BERT involves two LayerNorm sublayers – one after the MHA sublayer and another after the FFN sublayer.

Original: LayerNorm involves the calculation of the mean µ and standard deviation σ statistics along the feature dimension of size K , followed by an affine transformation:
Revised: LayerNorm involves the calculation of the mean µ and standard deviation σ statistics along the feature dimension of size K , followed by an affine transformation:

Original: [Equation]
Revised: [Equation]

Original: This remains true even during inference and demands the need to wait for all features of an input token to be processed before these statistics can be computed.
Revised: This remains true even during inference and demands the need to wait for all features of an input token to be processed before these statistics can be computed.

Original: It would be ideal if some form of batch statistics can be used as an alternative to normalize the inputs.
Revised: It would be ideal if some form of batch statistics can be used as an alternative to normalize the inputs.

Original: Doing so will eliminate the need to compute these statistics and present us the opportunity to fuse with adjacent linear layers in
Revised: Doing so will eliminate the need to compute these statistics and present us the opportunity to fuse with adjacent linear layers in

Original: BERT during inference.
Revised: BERT during inference.

Original: Recently, Shen et al.
Revised: Recently, Shen et al.

Original: (2020) proposed PowerNorm as a replacement for LayerNorm in Transformers and demonstrated that PowerNorm can outperform LayerNorm in neural machine translation and language modelling tasks.
Revised: (2020) proposed PowerNorm as a replacement for LayerNorm in Transformers and demonstrated that PowerNorm can outperform LayerNorm in neural machine translation and language modelling tasks.

Original: Similar to BatchNorm, PowerNorm normalizes the inputs along the batch dimension of size B and uses running statistics during inference.
Revised: Similar to BatchNorm, PowerNorm normalizes the inputs along the batch dimension of size B and uses running statistics during inference.

Original: However, it enforces only unit quadratic mean:
Revised: However, it enforces only unit quadratic mean:

Original: [Equation]
Revised: [Equation]

Original: During inference, the parameters γ , ψ , and β become constants which can be merged with the parameters in the adjacent linear layers.
Revised: During inference, the parameters γ , ψ , and β become constants which can be merged with the parameters in the adjacent linear layers.

Original: Specifically, if we were to replace LayerNorm with PowerNorm, the PowerNorm after the MHA sublayer can be fused with either the linear projection after the concatenation operation in Eq. 1 or the linear projection from the input layer to the intermediate layer of the following FFN.
Revised: Specifically, if we were to replace LayerNorm with PowerNorm, the PowerNorm after the MHA sublayer can be fused with either the linear projection after the concatenation operation in Eq. 1 or the linear projection from the input layer to the intermediate layer of the following FFN.

Original: Similarly, the PowerNorm after the FFN sublayer can be fused with either the linear projection from the intermediate layer to the output layer of the FFN or with the computation of the key, query, and value matrices in the next encoder layer.
Revised: Similarly, the PowerNorm after the FFN sublayer can be fused with either the linear projection from the intermediate layer to the output layer of the FFN or with the computation of the key, query, and value matrices in the next encoder layer.

Original: That said, due to the post-normalization architecture of BERT, it is not possible to do a complete fusion.
Revised: That said, due to the post-normalization architecture of BERT, it is not possible to do a complete fusion.

Original: Nonetheless, we can still reduce the two PowerNorm operations in every MA-BERT encoder layer to a single scaling operation during inference, which offers us considerable reduction in computation.
Revised: Nonetheless, we can still reduce the two PowerNorm operations in every MA-BERT encoder layer to a single scaling operation during inference, which offers us considerable reduction in computation.

Original: Given the prospect of minimizing the overhead of the normalization operation during inference, we propose to substitute LayerNorm in BERT with PowerNorm.
Revised: Given the prospect of minimizing the overhead of the normalization operation during inference, we propose to substitute LayerNorm in BERT with PowerNorm.

Original: Note that, unlike their work, we did not include a layer-scale operation before PowerNorm to prevent any additional overhead from being introduced.
Revised: Note that, unlike their work, we did not include a layer-scale operation before PowerNorm to prevent any additional overhead from being introduced.

Original: 3. K NOWLEDGE T RANSFER
Revised: 3. K NOWLEDGE T RANSFER

Original: With the three optimization techniques proposed in the previous subsections, we successfully achieve a BERT model with only matrix arithmetic and trivial ReLU operations.
Revised: With the three optimization techniques proposed in the previous subsections, we successfully achieve a BERT model with only matrix arithmetic and trivial ReLU operations.

Original: However, as a result of the modifications, a large accuracy gap exists between the untrained MA-BERT and the baseline.
Revised: However, as a result of the modifications, a large accuracy gap exists between the untrained MA-BERT and the baseline.

Original: In order to bridge this gap, we propose to leverage knowledge transfer.
Revised: In order to bridge this gap, we propose to leverage knowledge transfer.

Original: Many studies have demonstrated the effectiveness of knowledge distillation (Hinton et al., 2015) in bridging the performance gap between a student model and a teacher model.
Revised: Many studies have demonstrated the effectiveness of knowledge distillation (Hinton et al., 2015) in bridging the performance gap between a student model and a teacher model.

Original: The student model is trained to mimic the behaviour of a teacher model using a well-designed cost function that takes into account of the differences between the student output and the teacher output.
Revised: The student model is trained to mimic the behaviour of a teacher model using a well-designed cost function that takes into account of the differences between the student output and the teacher output.

Original: Unlike many studies which used knowledge distillation to compress BERT into smaller models, we leverage knowledge transfer with the intention to train a model of the same depth and hidden dimensions with a more efficient and hardware-friendly architecture.
Revised: Unlike many studies which used knowledge distillation to compress BERT into smaller models, we leverage knowledge transfer with the intention to train a model of the same depth and hidden dimensions with a more efficient and hardware-friendly architecture.

Original: Note that we are using the term knowledge transfer rather than knowledge distillation, given that the latter is often associated with a smaller model learning from a larger model.
Revised: Note that we are using the term knowledge transfer rather than knowledge distillation, given that the latter is often associated with a smaller model learning from a larger model.

Original: The following describes our procedure to transfer knowledge from an existing pretrained BERT base model to our modified model, MA-BERT base .
Revised: The following describes our procedure to transfer knowledge from an existing pretrained BERT base model to our modified model, MA-BERT base .

Original: 3.4. T EACHER M ODEL
Revised: 3.4. T EACHER M ODEL

Original: We used the pretrained BERT base model from HuggingFace as our teacher model.
Revised: We used the pretrained BERT base model from HuggingFace as our teacher model.

Original: All parameters of the teacher model were frozen to prevent the student model from chasing a moving target.
Revised: All parameters of the teacher model were frozen to prevent the student model from chasing a moving target.

Original: 3.4. S TUDENT M ODEL
Revised: 3.4. S TUDENT M ODEL

Original: For the student model (MA-BERT base ), we aggregated and implemented the three modifications described in the previous subsections on a BERT base model and initialized it with its pretrained parameters, wherever possible.
Revised: For the student model (MA-BERT base ), we aggregated and implemented the three modifications described in the previous subsections on a BERT base model and initialized it with its pretrained parameters, wherever possible.

Original: Furthermore, the newly added 2-layer neural network was pretrained in advance to approximate the softmax function.
Revised: Furthermore, the newly added 2-layer neural network was pretrained in advance to approximate the softmax function.

Original: For simplicity, the hidden layer size of the 2-layer neural network was set to be the same as the input sequence length.
Revised: For simplicity, the hidden layer size of the 2-layer neural network was set to be the same as the input sequence length.

Original: The proposed modifications were only applied on the embedding and encoder layers of BERT, and all parameters in these layers were left unfrozen, except for the position and token type embeddings.
Revised: The proposed modifications were only applied on the embedding and encoder layers of BERT, and all parameters in these layers were left unfrozen, except for the position and token type embeddings.

Original: 3.4. K NOWLEDGE T RANSFER O BJECTIVE
Revised: 3.4. K NOWLEDGE T RANSFER O BJECTIVE

Original: In order for the student to replicate the teacher’s behaviour as closely as possible, we adopted a similar layer-wise knowledge transfer technique proposed by Jiao et al.
Revised: In order for the student to replicate the teacher’s behaviour as closely as possible, we adopted a similar layer-wise knowledge transfer technique proposed by Jiao et al.

Original: First of all, we used the cross entropy (CE) between the student’s logits z S and the actual labels y to penalize any wrong classification:
Revised: First of all, we used the cross entropy (CE) between the student’s logits z S and the actual labels y to penalize any wrong classification:

Original: [Equation]
Revised: [Equation]

Original: Next, we considered the cross entropy between the student’s logits and the teacher’s logits, z T :
Revised: Next, we considered the cross entropy between the student’s logits and the teacher’s logits, z T :

Original: [Equation]
Revised: [Equation]

Original: Compared to the actual labels, the teacher’s logits provide much more information for the student to learn.
Revised: Compared to the actual labels, the teacher’s logits provide much more information for the student to learn.

Original: A higher temperature t smoothens the probability distribution, which allows the student to better learn the intricacies of the teacher model.
Revised: A higher temperature t smoothens the probability distribution, which allows the student to better learn the intricacies of the teacher model.

Original: Thirdly, to ensure each layer of our student model, including the embedding layer, learns from the corresponding layer of teacher model, we included the following mean-squared error (MSE) objective:
Revised: Thirdly, to ensure each layer of our student model, including the embedding layer, learns from the corresponding layer of teacher model, we included the following mean-squared error (MSE) objective:

Original: [Equation]
Revised: [Equation]

Original: H Si and H Ti are the hidden states of the student and teacher models, respectively, and l is the number of layers in the model.
Revised: H Si and H Ti are the hidden states of the student and teacher models, respectively, and l is the number of layers in the model.

Original: Lastly, given that we are replacing the softmax function in the MHA sublayer, we also penalized the discrepancy between the student’s attention outputs
Revised: Lastly, given that we are replacing the softmax function in the MHA sublayer, we also penalized the discrepancy between the student’s attention outputs

Original: A Si and teacher’s attention outputs
Revised: A Si and teacher’s attention outputs

Original: A Ti using:
Revised: A Ti using:

Original: [Equation]
Revised: [Equation]

Original: Altogether, the overall objective function for knowledge transfer can be summarized as follows:
Revised: Altogether, the overall objective function for knowledge transfer can be summarized as follows:

Original: [Equation]
Revised: [Equation]

Original: α is used to balance the importance of following the actual hard target and the teacher’s soft target, β is used to balance the importance of mimicking the teacher’s internal layer outputs relative to mimicking the final output, and γ is a factor to scale L attn to a similar magnitude as L hidn .
Revised: α is used to balance the importance of following the actual hard target and the teacher’s soft target, β is used to balance the importance of mimicking the teacher’s internal layer outputs relative to mimicking the final output, and γ is a factor to scale L attn to a similar magnitude as L hidn .

Original: With this overall objective, the student is compelled to not only mimic the final output of the teacher, but also from its hidden layer and attention outputs.
Revised: With this overall objective, the student is compelled to not only mimic the final output of the teacher, but also from its hidden layer and attention outputs.

Original: Throughout our study, we kept t = 15 , α = 0 .
Revised: Throughout our study, we kept t = 15 , α = 0 .

Original: 9 , β = 1 , and γ = 100 .
Revised: 9 , β = 1 , and γ = 100 .

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we detail our pretraining setup and evaluate the performance of MA-BERT base on several downstream tasks.
Revised: In this section, we detail our pretraining setup and evaluate the performance of MA-BERT base on several downstream tasks.

Original: We also describe the details of additional studies that we have conducted which attempt to highlight and exploit the benefits of our approach.
Revised: We also describe the details of additional studies that we have conducted which attempt to highlight and exploit the benefits of our approach.

Original: Throughout this section, all results pertain to a model input sequence length of 128 tokens.
Revised: Throughout this section, all results pertain to a model input sequence length of 128 tokens.

Original: 4. P RETRAINING S ETUP
Revised: 4. P RETRAINING S ETUP

Original: We pretrained MA-BERT base via knowledge transfer on a concatenation of BookCorpus (Zhu et al., 2015) and English Wikipedia using only the Masked Language Modelling objective (Devlin et al., 2018) with masking probability of 15%.
Revised: We pretrained MA-BERT base via knowledge transfer on a concatenation of BookCorpus (Zhu et al., 2015) and English Wikipedia using only the Masked Language Modelling objective (Devlin et al., 2018) with masking probability of 15%.

Original: We used AdamW optimizer with learning rate of 2 e − 5 , 10,000 warm-up steps from a base learning rate of 5 e − 7 , and a linear decay of learning rate.
Revised: We used AdamW optimizer with learning rate of 2 e − 5 , 10,000 warm-up steps from a base learning rate of 5 e − 7 , and a linear decay of learning rate.

Original: Training was done with a batch size of 256 over 3 epochs, with each epoch using a differently masked dataset, on a single NVIDIA GeForce RTX 3090 for a total of 3.5 days.
Revised: Training was done with a batch size of 256 over 3 epochs, with each epoch using a differently masked dataset, on a single NVIDIA GeForce RTX 3090 for a total of 3.5 days.

Original: 4. E VALUATION ON D OWNSTREAM T ASKS
Revised: 4. E VALUATION ON D OWNSTREAM T ASKS

Original: To evaluate the performance of MA-BERT base , we used the GLUE benchmark (Wang et al., 2018), which consists of nine natural language understanding tasks, and the IMDb sentiment classification task (Maas et al., 2011).
Revised: To evaluate the performance of MA-BERT base , we used the GLUE benchmark (Wang et al., 2018), which consists of nine natural language understanding tasks, and the IMDb sentiment classification task (Maas et al., 2011).

Original: Following many studies, we excluded WNLI given its small dataset.
Revised: Following many studies, we excluded WNLI given its small dataset.

Original: For each task, we applied task-specific knowledge transfer by having MA-BERT base learn from a finetuned BERT base .
Revised: For each task, we applied task-specific knowledge transfer by having MA-BERT base learn from a finetuned BERT base .

Original: The same knowledge transfer objective and hyperparameter settings outlined in section 3.4.3 were used.
Revised: The same knowledge transfer objective and hyperparameter settings outlined in section 3.4.3 were used.

Original: As our goal is to replicate the performance of the teacher, our baseline for comparison is simply the teacher, which is BERT base in this case.
Revised: As our goal is to replicate the performance of the teacher, our baseline for comparison is simply the teacher, which is BERT base in this case.

Original: For smaller datasets (CoLA, MRPC, RTE and STS-B), we used a batch size of 16 and fine-tuned for 10 epochs, and for the remaining datasets, we used a batch size of 32 and fine-tuned for 5 epochs.
Revised: For smaller datasets (CoLA, MRPC, RTE and STS-B), we used a batch size of 16 and fine-tuned for 10 epochs, and for the remaining datasets, we used a batch size of 32 and fine-tuned for 5 epochs.

Original: The same learning rate of 2 e − 5 was used across all tasks and the sequence length was fixed at 128.
Revised: The same learning rate of 2 e − 5 was used across all tasks and the sequence length was fixed at 128.

Original: To account for the variance of different runs, we report the median of 5 runs with different random initialization seeds for each task.
Revised: To account for the variance of different runs, we report the median of 5 runs with different random initialization seeds for each task.

Original: The results are shown in Table 2 and 3.
Revised: The results are shown in Table 2 and 3.

Original: Note that the teacher with the best performance among the 5 runs was used to finetune the student during taskspecific knowledge transfer.
Revised: Note that the teacher with the best performance among the 5 runs was used to finetune the student during taskspecific knowledge transfer.

Original: The best performance of both the teacher and student is reported under
Revised: The best performance of both the teacher and student is reported under

Original: Appendix A.1.
Revised: Appendix A.1.

Original: On average, despite the modifications, MA-BERT base was able to achieve similar performance on the downstream tasks as compared to BERT base .
Revised: On average, despite the modifications, MA-BERT base was able to achieve similar performance on the downstream tasks as compared to BERT base .

Original: In fact, for several tasks, MA-BERT base managed to outperform its teacher.
Revised: In fact, for several tasks, MA-BERT base managed to outperform its teacher.

Original: This demonstrates that adopting our approach does not lead to a degradation in the average performance of the model.
Revised: This demonstrates that adopting our approach does not lead to a degradation in the average performance of the model.

Original: 4. H YPERPARAMETER T UNING OF α AND β
Revised: 4. H YPERPARAMETER T UNING OF α AND β

Original: Prior to the collection of results in Table 2 and 3, we conducted a simple hyperparameter study to assess the impact of our choices for α and β in our knowledge transfer objective during task-specific knowledge transfer.
Revised: Prior to the collection of results in Table 2 and 3, we conducted a simple hyperparameter study to assess the impact of our choices for α and β in our knowledge transfer objective during task-specific knowledge transfer.

Original: We varied α and β , which effectively control how strictly the student should replicate the teacher’s final and hidden layer outputs, and evaluated MA-BERT base on the 4 smaller datasets of the GLUE benchmark.
Revised: We varied α and β , which effectively control how strictly the student should replicate the teacher’s final and hidden layer outputs, and evaluated MA-BERT base on the 4 smaller datasets of the GLUE benchmark.

Original: The results are shown in Table 4.
Revised: The results are shown in Table 4.

Original: Empirically, we found that MABERT base performs the best on downstream tasks under the strict guidance of the finetuned BERT base model with the optimal hyperparameter values of α = 0 .
Revised: Empirically, we found that MABERT base performs the best on downstream tasks under the strict guidance of the finetuned BERT base model with the optimal hyperparameter values of α = 0 .

Original: 9 and β =
Revised: 9 and β =

Original: [Table] [Table]
Revised: [Table] [Table]

Original: selected GLUE tasks.
Revised: selected GLUE tasks.

Original: [Table]
Revised: [Table]

Original: 4. S HARING S OFTMAX A PPROXIMATION N EURAL N ETWORK A CROSS L AYERS
Revised: 4. S HARING S OFTMAX A PPROXIMATION N EURAL N ETWORK A CROSS L AYERS

Original: In the previous subsection, the softmax function for each encoder layer of MA-BERT base were approximated using a different 2-layer neural network.
Revised: In the previous subsection [Table] the softmax function in each encoder layer of MA-BERT base were approximated using a different 2-layer neural network.

Original: However, given that it is the same softmax function, it would make sense if we can share the same 2-layer neural network across all layers.
Revised: However, given that it is the same softmax function [Table] it would make sense if we can share the same 2-layer neural network across all layers.

Original: This would  help us reduce the additional parameters that our approach introduced.
Revised: Doing so would also help us reduce the additional parameters introduced by our approach. In fact [Table]

Revised: we found that the weights learnt by the 2-layer neural network for each encoder layer in MA-BERT base are largely similar.

Revised: A visualization of the weights learnt can be found in Figure 5 under Appendix

Original:  We attempted to retrain MA-BERT base with parameter sharing in the softmax approximation neural network using the same training procedure.
Revised: A.2. We attempted to retrain MA-BERT base with parameter sharing in the softmax approximation neural network using the same training procedure.

Original: Empirically, we found that sharing the 2-layer neural network to approximate softmax preserves the performance of the modified model.
Revised: Empirically, we found that sharing the 2-layer neural network to approximate softmax preserves the performance of the modified model.

Original: 4. V ERSATILITY S TUDY
Revised: 4. V ERSATILITY S TUDY

Original: To evaluate the versatility of our approach
Revised: To evaluate the versatility of our approach

Original: [Table]
Revised: [Table]

Original: we also attempted our modifications and training procedure on HuggingFace’s distilBERT (Sanh et al., 2019) and distilRoBERTa to produce MAdistilBERT and MA-distilRoBERTa, respectively.
Revised: we also attempted our modifications and training procedure on HuggingFace’s DistilBERT (Sanh et al., 2019) and DistilRoBERTa to produce MADistilBERT and MA-DistilRoBERTa, respectively.

Original: Following distilRoBERTa’s training procedure, MA-distilRoBERTA was insteadpretrained on the OpenWebText (Gokaslan & Cohen, 2019) dataset.
Revised: Following DistilRoBERTa’s training procedure, MA-DistilRoBERTA was instead dataset.

Original: The remaining training procedure was kept unchanged.
Revised: The remaining training procedure was kept unchanged.

Original: We observed that both MAdistilBERT and MA-distilRoBERTa suffered a significant drop in performance on the CoLA task.
Revised: We observed that both MADistilBERT and MA-DistilRoBERTa suffered a significant drop in performance on the CoLA task.

Original: Nonethless, in terms of average performance, MA-distilBERT managed to slightly outperform its teacher, distilBERT.
Revised: Nonethless, in terms of average performance, MA-DistilBERT managed to slightly outperform its teacher, DistilBERT.

Original: On the other hand, although MA-distilRoBERTa was not able to achieve similar average performance as distilRoBERTa, it was still able to outperform distilBERT, which has the same hidden dimensions and depth.
Revised: On the other hand, although MA-DistilRoBERTa was not able to achieve similar average performance as DistilRoBERTa, it was still able to outperform DistilBERT, which has the same hidden dimensions and depth.

Original: This versatility study also illustrates that our approach is orthogonal to model compression via knowledge distillation.
Revised: This versatility study also illustrates that our approach is orthogonal to model compression via knowledge distillation.

Original: We can first distil a model into a shallower version, then apply our modifications to make it even more efficient and hardware friendly for inference.
Revised: We can first distil a model into a shallower version, then apply our modifications to make it even more efficient and hardware friendly for inference.

Original: 4. A DOPTING L INEAR S ELF -A TTENTION
Revised: 4. A DOPTING L INEAR S ELF -A TTENTION

Original: With the use of a neural network to approximate the softmax operation, we have inevitably introduced additional matrix multiplication operations.
Revised: With the use of a neural network to approximate the softmax operation, we have inevitably introduced additional matrix multiplication operations.

Original: The increase in matrix multiplication scales in
Revised: The increase in matrix multiplication scales in

Original: O ( s 3 ) if the hidden layer size of the neural network follows the input sequence length s .
Revised: O ( s 3 ) if the hidden layer size of the neural network follows the input sequence length s .

Original: That said, for relatively short sequence lengths, the increase is marginal.
Revised: That said, for relatively short sequence lengths, the increase is marginal.

Original: For s = 128 , the increase in the total number of matrix multiplication compared to the baseline is around 5.4% (Figure 2).
Revised: For s = 128 , the increase in the total number of matrix multiplication compared to the baseline is around 5.4% (Figure 2).

Original: Nonetheless, the issue of scalability comes to light.
Revised: Nonetheless, the issue of scalability comes to light.

Original: For longer sequence lengths, we suggest a possible solution that can reduce this overhead.
Revised: For longer sequence lengths, we suggest a possible solution that can reduce this overhead.

Original: Recently,
Revised: Recently,

Original: Wang et al.
Revised: Wang et al.

Original: (2020) proposed self-attention with linear complexity.
Revised: (2020) proposed self-attention with linear complexity.

Original: This is achieved by projecting the key and value matrices from dimensions R s × d into smaller dimensions of R k × d using the projection matrices E, F ∈ R k × s and setting k ≪ s : 
Revised: This is achieved by projecting the key and value matrices from dimensions R s × d into smaller dimensions of R k × d using the projection matrices E, F ∈ R k × s and setting k ≪ s : [Equation]

Original: [Equation]
Revised: [Equation]

Original: Here, d denotes the hidden dimension of the model while k denotes the new projected dimension.
Revised: Here, d denotes the hidden dimension of the model while k denotes the new projected dimension.

Original: By adopting linear self-attention, softmax has to be computed only along the reduced dimension k , which enables us to use a smaller 2-layer neural network to replace softmax.
Revised: By adopting linear self-attention, softmax has to be computed only along the reduced dimension k , which enables us to use a smaller 2-layer neural network to replace softmax.

Original: This not only helps to reduce the overheads introduced by the additional neural network to approximate softmax, but also allows us to leverage the speedup offered by the linear self-attention.
Revised: This not only helps to reduce the overheads introduced by the additional neural network to approximate softmax, but also allows us to leverage the speedup offered by the linear self-attention.

Original: For instance, by letting k = 0 .
Revised: For instance, by letting k = 0 .

Original: 5 s , significant savings can be achieved as illustrated in Figure 2.
Revised: 5 s , significant savings can be achieved as illustrated in Figure 2.

Original: To take this even further, we can also fuse the linear operations in the softmax approximation neural network with both the projection of the key and value matrices during inference, hiding the overhead introduced.
Revised: To take this even further, we can also fuse the linear operations in the softmax approximation neural network with both the projection of the key and value matrices during inference, hiding the overhead introduced.

Original: To test this out, we trained MA-DistilRoBERTa with linear attention by setting k = 64 using the same procedure.
Revised: To test this out, we trained MA-DistilRoBERTa with linear attention by setting k = 64 using the same procedure.

Original: The loss term L attn was removed from the knowledge transfer objective due to the different attention output dimensions of the modified model and DistilRoBERTa.
Revised: The loss term L attn was removed from the knowledge transfer objective due to the different attention output dimensions of the modified model and DistilRoBERTa.

Original: The results are reported in Table 2 and 3.
Revised: The results are reported in Table 2 and 3.

Original: We noticed a 3.3% drop in performance of the modified model compared with the baseline DistilRoBERTa.
Revised: We noticed a 3.3% drop in performance of the modified model compared with the baseline DistilRoBERTa.

Original: Nonetheless, it was still able to achieve comparable average performance (0.37% drop) as baseline DistilBERT.
Revised: Nonetheless, it was still able to achieve comparable average performance (0.37% drop) as baseline DistilBERT.

Original: A revision on the training procedure might be necessary to further improve its performance and we will leave this to future studies.
Revised: A revision on the training procedure might be necessary to further improve its performance and we will leave this to future studies.

Original: [Figure]
Revised: [Figure]

Original: 4. E VALUATION ON C OMPUTATIONAL E FFICIENCY
Revised: 4. E VALUATION ON C OMPUTATIONAL E FFICIENCY

Original: To evaluate the time efficiency of MA-BERT during inference, we implemented a single encoder layer of BERT and MA-BERT in C++ using the Eigen3 library.
Revised: To evaluate the time efficiency of MA-BERT during inference, we implemented DistilBERT and MA-DistilBERT on three different hardware, namely, i7-1260P CPU, RTX 3090 GPU, and Google

Original: All weights, biases, and calculations were in 32-bit floating-point precision. We ran the code recursively for 12 times, simulating an inference made by an encoder block of a BERT base model, and measured the time taken forsuch consecutive inferences.
Revised: TPUv2.8.

Original: We then computed the average time taken for a single  layer inference.The C++ code was executed on both Intel’s i7-6500U and i7-1260P CPUs and was compiled using TDM-GCC compiler with various levels of optimization.
Revised: We used C++ with Eigen3 library for CPU implementation, Cuda C++ with Thrust and cuBLAS libraries for GPU implementation, and PyTorch for TPU implementation. Using the IMDb test dataset, we measured the average time taken for a single encoder layer inference.

Original: To account for the variance in the time taken, we report the median of 5 different runs.
Revised: To account for the variance in the time taken, we report the median of 5 different runs.

Original: The results are summarized in Table 5.
Revised: The results are summarized in Table 5.

Original: Despite the higher number of matrix multiplications involved, MA-BERT is able to offer an inference time speedupof as much as 1.27 times compared to BERT on general CPUs.
Revised: Despite the higher number of matrix multiplications involved, MA-DistilBERT is able to offer an inference time speedup on various hardware. Note that the speedups reported also applies for MA-BERT base and MA-DistilRoBERTa, given that their encoder layers are all identical.

Revised: This demonstrates that our approach not only eliminates the need to compute any complex non-linear functions, but also helps to reduce the inference time.

Original: [Table]
Revised: [Table]

Original:  A BLATION S TUDY
Revised:  A BLATION S TUDY

Original: In this last section, we highlight the importance of having the baseline teacher’s guidance during
Revised: In this last section, we highlight the importance of having the baseline teacher’s guidance during

Original: MA-BERT’s finetuning phase.
Revised: MA-BERT’s finetuning phase.

Original: 5. R EMOVAL OF T EACHER ’ S G UIDANCE
Revised: 5. R EMOVAL OF T EACHER ’ S G UIDANCE

Original: To demonstrate the importance of task-specific knowledge transfer during finetuning, we repeated our finetuning procedure on the GLUE tasks by replacing the knowledge transfer objective with only the hard target loss for our student model to optimize.
Revised: To demonstrate the importance of task-specific knowledge transfer during finetuning, we repeated our finetuning procedure on the GLUE tasks by replacing the knowledge transfer objective with only the hard target loss for our student model to optimize.

Original: The results are reported in Table 6.
Revised: The results are reported in Table 6.

Original: We found that without any guidance from the teacher, the finetuning process is often unstable with the student model ending up performing worse in all tasks.
Revised: We found that without any guidance from the teacher, the finetuning process is often unstable with the student model ending up performing worse in all tasks.

Original: Clearly, knowledge transfer plays a significant role in the success of our approach.
Revised: Clearly, knowledge transfer plays a significant role in the success of our approach.

Original: [Table]
Revised: [Table]

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper
Revised: In this paper

Original: [Table]
Revised: [Table]

Original: we introduced MA-BERT, an efficient and hardware-friendly version of BERT with all its complex non-linear functions eliminated.
Revised: we introduced MA-BERT [Table] an efficient and hardware-friendly version of BERT with all its complex non-linear functions eliminated.

Original: Through knowledge transfer during pretraining and finetuning, MA-BERT was able to attain comparable performance as BERT.
Revised: Through knowledge transfer during pretraining and finetuning, MA-BERT is able to attain comparable performance as BERT.

Original: Yet [Table] compared to BERT,
Revised: Yet, compared to BERT,

Original: MA-BERT offers as much as 27% reduction in inference time with potential savings in hardware resources [Table] making it a suitable solution for applications on edge devices.
Revised: MA-BERT offers a reduction in inference time with potential savings in hardware resources, making it a suitable solution for applications on edge devices.

Original: In the future
Revised: In the future

Original: [Table]
Revised: [Table]

Original: it would be interesting to study the applicability of our approach on other transformer-based models
Revised: it would be interesting to study the applicability of our approach on other transformer-based models.

Original: [Table]

Original: A A PPENDIX A. P ERFORMANCE OF B EST P ERFORMING T EACHER AND S TUDENT M ODELS
Revised: A A PPENDIX A. P ERFORMANCE OF B EST P ERFORMING T EACHER AND S TUDENT M ODELS

Original: During task-specific knowledge transfer, the best-performing finetuned teacher was used as the teacher for the pretrained student model.
Revised: During task-specific knowledge transfer, the best-performing finetuned teacher was used as the teacher for the pretrained student model.

Original: Table 7 shows the best performance achieved by the baseline teacher models and the modified student models among the 5 runs for each task.
Revised: Table 7 shows the best performance achieved by the baseline teacher models and the modified student models among the 5 runs for each task.

Original: [Table]
Revised: [Table]

Original: ) ( ▼ 0.5)
Revised: ) ( ▼ 0.5)

Revised: A. I NTERESTING F INDINGS ON S OFTMAX A PPROXIMATION

Revised: After pretraining MA-BERT base , we decided to evaluate the effectiveness of the 2-layer neural network in approximating softmax.

Revised: We fed MA-BERT base with the sample input sequence “the quick brown fox jumps over the lazy dog” and plotted out the multi-head self-attention scores produced in each of its encoder layer.

Revised: We then compared them with the corresponding encoder layer attention scores in BERT base .

Revised: Figure 3 shows a side-by-side comparison of the attention scores produced in the 4 th , 8 th , and 12 th encoder layers of MA-BERT base and BERT base .

Revised: We observe that the attention scores produced by both models are similar, which demonstrates the effectiveness of the 2-layer neural network in approximating softmax.

Revised: The same observation is made when we compare the attention outputs in MA-DistilBERT and DistilBERT in Figure 4.

Revised: Moreover, we decided to look into the weights that are learnt by the 2-layer neural network to investigate if any meaningful mappings are learnt.

Revised: Figure 5 shows the learnt weights mapping the input layer to the hidden layer and the hidden layer to the output layer for each encoder layer.

Revised: Interestingly, despite giving MA-BERT base the opportunity to learn a different mapping function to approximate softmax in different encoder layers, each 2-layer neural network ended up learning a similar mapping.

Revised: In fact, when we added the constraint of sharing the 2-layer neural network across all 12 encoder layers in MA-BERT base (Shared Softmax), a similar mapping is also learnt.

Revised: Note that the mappings are rather sparse and do not seem to be random, which suggests that the 2-layer neural network has learnt something meaningful.

Revised: The same observation is made when we looked into the learnt weights in MA-DistilBERT in Figure 6.

Revised: In fact, we can see that the mappings learnt in MA-DistilBERT and MA-BERT base are similar as well.

Revised: This further suggests that there is some underlying function that can approximate the role of softmax in self-attention and we believe that this will serve as an interesting point for future research

Revised: [Table]

Revised: [Table]

Revised: l a z y d

Revised: o g

Revised: The quick brown fox jumps over the lazy dog

Revised: Head

Revised: T h e q u

Revised: i c k b r o w n f o x j u m p s o v e r t h e l a z y d o g

Revised: The quick brown fox jumps over the lazy dog

Revised: [Figure]

Revised: T h e q u

Revised: i c k b r o w n f o x j u m p s o v e r t h e l a z y d o g

Revised: The quick brown fox jumps over the lazy dog

Revised: [Figure]

Revised: T h e q u

Revised: i c k b r o w n f o x j u m p s o v e r t h e l a z y d o g

Revised: The quick brown fox jumps over the lazy

Revised: [Figure]

Revised: T h e q u

Revised: i c k b r o w n f o x j u m p s o v e r t h e l a z y d o g

Revised: The quick

Revised: [Figure]

Revised: jumps over

Revised: [Figure]

Revised: lazy

Revised: [Figure]

Revised: T h e q u

Revised: i c k b r o w n f o x j u m p s o v e r t h e l a z y d o g

Revised: [Figure]

Revised: [Table]

Revised: jumps over the lazy dog

Revised: Head

Revised: T h e q u

Revised: i c k b r o w n f o x j u m p s o v e r t h e l a z y d o g

Revised: The quick brown fox jumps over the lazy dog

Revised: Head

Revised: T h e q u

Revised: i c k b r o w n f o x j u m p s o v e r t h e l a z y d o g

Revised: The quick brown fox jumps over the lazy dog

Revised: Head

Revised: T h e q u

Revised: i c k b r o w n f o x j u m p s o v e r t h e l a z y d o g

Revised: The quick brown fox jumps over the lazy dog

Revised: [Figure]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Figure]