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.