A Novel Weight Initialization Method for Neural Networks with Nonlinear Output

In this paper, we propose a weight initialization method to address the challenges associated with initializing models that utilize the softmax function to convert their final output into probabilities. Utilizing the first-order Taylor expansion, we approximate the model's nonlinear output with a linear function, which enables us to fine-tune the weights of the last layer by minimizing the initial loss function. For the other layers, we apply the well-established Kaiming initialization technique to ensure that the variance of the output is equal to the variance of the input, ensuring a more stable training process. This strategic optimization yields weights that are close to optimal, thereby accelerating the convergence rate and enhancing the efficiency of the training process. Experiments comparing our method against four traditional initialization methods, showing our approach has a significant advantage, especially during the initial training phase, while remaining competitive in the later stages of training.