# MicroBERT: Distilling MoE-Based Knowledge from BERT into a Lighter Model
## 3 Methods
### Distillation Process
The following objectives have been defined:  
\[\mathcal{L}_{\text{hidn}}=\text{MSE}\Big{(}\mathbf{H}^{T},\mathbf{H}^{S}\mathbf{W}_{h} \Big{)} \tag{2}\]  
where \(\mathbf{H}^{S}\in\mathbb{R}^{l\times d^{\prime}}\) and \(\mathbf{H}^{T}\in\mathbb{R}^{l\times d}\) are matrices that represent the hidden states of the student and teacher models, as defined by Equation (2). Thus the hidden sizes of the classroom model and student model are defined as \(d\) and \(d^{\prime}\). We want to maximize the compression of the model, so the student model \(d^{\prime}\) is chosen to be frequently smaller than \(d\). Also, the hidden state of the student model is transformed to the same state dimension as the teacher model by the learnable linear transformation matrix mapping \(\mathbf{W}_{h}\in\mathbb{R}^{d^{\prime}\times d}\). **Embedding Layer Distillation**. We also performed distillation on the embedding layer, aiming to achieve similar goals as the hidden-state-based distillation:  
\[\mathcal{L}_{\text{embd}}=\text{MSE}\Big{(}\mathbf{E}^{T},\mathbf{E}^{S}\mathbf{W}_{e} \Big{)} \tag{3}\]  
where the matrices \(\mathbf{E}^{S}\) and \(\mathbf{E}^{T}\) stand for the networks' respective embeddings in students' and teachers' networks. They share the same structure as the hidden state matrices in this study. An analogous linear transformation to \(\mathbf{W}_{h}\) is the matrix \(\mathbf{W}_{e}\). **MoE-guided distillation**. We use the MoE proposed by [PERSON] [30], where the MoE layers for the Transformers consist of \(E\) Prediction Layer, such that (\(PL_{1}\dots\ PL_{E}\)). \[PL_{e}(x_{s})=wo_{e}Softmax(wi_{e}x_{s}) \tag{4}\]  
\[y_{s}=\sum_{e=1}^{E}G_{s,e}PL_{e}(x_{s}) \tag{5}\]  
The input token for the MoE layer at position s is \(x_{s}\), and each \(PL_{e}\) denotes a one-layer neural network with a Softmax activation function. The projection weights for the e-th expert's input and output are \(wi_{e}\) and \(wo_{e}\). Finally, \(G_{s,e}\) is the vector produced by the gating network, commonly referred to as the router. For each expert, the vector's values are mostly zeros, with the exception of one positive value. Using this vector, we route the token to a restricted set of experts. The items chosen from \(G_{s}\), \(E\), determine the expert's contribution to the final output \(y_{s}\). It should be noted that in order to be comparable with the previous study, we selected the top 1 weight experts for each case in this work. The gating network \(G_{s}\), \(E\), where (1) expert utilization must be balanced and (2) the function must be efficient to execute at scale, must be carefully examined for efficiency reasons [30]. **Prediction Layer Distillation**. For the prediction layer component of our model, we employ the knowledge-distillation method proposed by [PERSON] et al. [31]. This method involves calculating the soft cross-entropy loss by taking the logarithm of the outputs from both the student model and the teacher model. Additionally, the soft labels generated by the trained teacher model can be transferred to any task. \[\mathcal{L}_{\text{pred}}\,=\text{CE}\Big{(}\boldsymbol{z}^{T}/t,\boldsymbol{ z}^{S}/t\Big{)} \tag{6}\]  
where \(CE\) stands for cross entropy loss, \(t\) stands for temperature value and \(\boldsymbol{z}^{S}\) and \(\boldsymbol{z}^{T}\) are the logits vectors predicted by the student and teacher, respectively. In our experiment, \(t=1\) shows good performance. Finally, the above Equations (2), (3) and (6) distillation losses are summarized to yield the following target losses:  
\[\mathcal{L}_{\text{layer}}\,=\,\mathcal{L}_{\text{embd}}\,+\mathcal{L}_{\text {hidn}}\,+\mathcal{L}_{\text{pred}} \tag{7}\]  
where calculate \(\mathcal{L}_{\text{embd}}\,\) loss if \(m=0\), \(\mathcal{L}_{\text{hidn}}\,\) loss if \(M\geq m>0\) and \(\mathcal{L}_{\text{pred}}\,\) loss if \(m=M+1\).