## 2 System Model
### Two-Stage KD  
In order to further compress the model, KD is employed to transfer the knowledge from a pre-trained RIMformer model (the teacher model) to a small model with a lightweight design and reduced layers (the student model). This process leverages the teacher's learned representations to guide the student's training, resulting in a model that approximates the teacher's performance with a reduced complexity.  
It has been observed that when the discrepancy between the teacher and student networks is too high, there is a decline in the efficiency of KD. To address this challenge, we propose a multi-step KD approach involving an intermediate-sized network (an assistant model) to bridge the gap between the student and teacher networks. The overall process is depicted in Figure 5. By introducing an intermediate assistant teacher model, we create a smoother knowledge transfer process. The intermediate assistant teacher model incorporates lightweight design improvements but retains a structure closer to the original teacher. It acts as a \"bridge\" by effectively distilling the critical features and representations from the teacher model in a simplified form. The student model can then learn from the assistant model in the second stage, with improved alignment due to the significantly reduced gap in model capacity and complexity. The two-stage approach strikes a balance between computational efficiency and performance. Although the addition of an intermediate model introduces an additional training step, this step is essential to help the lightweight student model achieve competitive performance in interference mitigation without sacrificing the benefits of the proposed lightweight designs.  
The distillation process integrates response-based knowledge and feature-based knowledge. Response-based knowledge refers to the neural responses of the teacher model's final output layer. The principal idea is to have the student model directly emulate the teacher model's final predictions. Data representations from the intermediate layers can also serve as knowledge for training the student model, also known as feature-based knowledge. In this paper, the L2 norm is utilized to measure the similarity between the feature maps and the output layer. The distillation stage loss function is represented as follows:  
\[\mathcal{L}_{\text{KD}}=\frac{1}{2}\Big{[}L_{\text{RMSE}}\Big{(}K_{\text{F}},\,\widetilde{K}_{\text{F}}\Big{)}+L_{\text{RMSE}}\Big{(}K_{\text{R}},\, \widetilde{K}_{\text{R}}\Big{)}\Big{]}. \tag{11}\]The overall training loss function is a weighted sum of the distillation loss and the actual value of the time-frequency mixed loss, which can be calculated as follows:  
\[\mathcal{L}=\alpha\times\mathcal{L}_{\text{KD}}+(1-\alpha)\times\mathcal{L}_{ \text{TF}}, \tag{12}\]  
where \(\alpha\) is a weight that balances the two parts of the loss.  
Furthermore, KD is optimized through dynamic loss balancing. As shown in Figure 6, the proposed training strategy involves adjusting the weight parameter in the network loss function throughout the training process. Initially, the integration of distilled knowledge from the teacher network is accentuated with a large weight, which fosters a swift convergence of the student network towards an optimal solution. Subsequently, a medium weight strikes a balance between the influences of KD and label prediction losses. The student network's comprehension is refined with mitigated excessive reliance on pre-existing knowledge. Lastly, the employment of a smaller weight in later stages alleviates the risk of overfitting the teacher's knowledge and facilitates generalization to unseen data. This dynamic methodology provides a nuanced and efficient approach to distill knowledge while ensuring the adaptability and resilience of the student network.  
Figure 5: Architecture diagram for two-stage KD.  
Figure 6: Optimizing KD through dynamic loss balancing.