OptiDistillNet: Learning nonlinear pulse propagation using the student-teacher model

We present a unique approach for learning the pulse evolution in a nonlinear fiber using a deep convolutional neural network (CNN) by solving the nonlinear Schrodinger equation (NLSE). Deep network model compression has become widespread for deploying such models in real-world applications. A knowledge distillation (KD) based framework for compressing a CNN is presented here. The student network, termed here as OptiDistillNet has better generalisation, has faster convergence, is faster and uses less number of trainable parameters. This work represents the first effort, to the best of our knowledge, that successfully applies a KD-based technique for any nonlinear optics application. Our tests show that even by reducing the model size by up to 91.2%, we can still achieve a mean square error (MSE) which is very close to the MSE of 1.04*10-5 achieved by the teacher model. The advantages of the suggested model include the use of a simple architecture, fast optimization, and improved accuracy, opening up applications in optical coherent communication systems. (c) 2022 Optica Publishing Group under the terms of the Optica Open Access Publishing Agreement

Automated summary

This paper presents a unique approach using a deep convolutional neural network to solve the nonlinear Schrodinger equation for learning the pulse evolution. A knowledge distillation-based technique is applied to compress the model. The results show that even with a reduction of up to 91.2% in model size, the achieved mean square error is very close to that of the teacher model.