Robust Weight Initialization for Tanh Neural Networks with Fixed Point Analysis

Authors: Hyun woo Lee, Hayoung Choi, Hyunju Kim
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Weight initialization, Signal propagation, Physics informed neural networks

Abstract:
As a neural network's depth increases, it can improve generalization performance. However, training deep networks is challenging due to gradient and signal propagation issues. To address these challenges, extensive theoretical research and various methods have been introduced. Despite these advances, effective weight initialization methods for tanh neural networks remain insufficiently investigated. This paper presents a novel weight initialization method for neural networks with tanh activation function. Based on an analysis of the fixed points of the function $\tanh(ax)$, the proposed method aims to determine values of $a$ that mitigate activation saturation. A series of experiments on various classification datasets and physics-informed neural networks demonstrates that the proposed method outperforms Xavier initialization methods (with or without normalization) in terms of robustness across different network sizes, data efficiency, and convergence speed. Code is available at https://github.com/1HyunwooLee/Tanh-Init.

Decision: accept (poster)
Reviewer scores: 8, 6, 5, 6

Reviews and discussion:
<thread>
creation date:
1729820121208

thread_type:
official_review

<thread_object>
creator: reviewer_ycax

<content>
summary:
The paper proposes a new method to initialize weights for FCNNs and PINNs with tanh activation. The paper claims that the proposed weight initialization method will not lead to diminishing activations for very deep networks unlike Xavier weight initialization. The paper claims that the proposed weight initialization is robust to network depth and number of units in hidden leayers.

strengths:
Originality
The paper presents a novel weight initialization method specifically designed for tanh-based neural networks, addressing an understudied area in neural network initialization. This approach is distinct in its use of fixed-point analysis to prevent activation saturation and improve training robustness across network sizes, particularly in the context of Physics-Informed Neural Networks (PINNs). By emphasizing robustness and performance consistency in both traditional classification tasks and PINNs, the paper makes a valuable contribution to the field of weight initialization. The originality is strong, given the lack of prior research focusing on tanh-based initialization methods.

Quality
The paper demonstrates high quality in both theoretical and experimental aspects. The method is grounded in rigorous mathematical analysis, leveraging fixed-point theory to derive conditions that ensure stable activation propagation. The provided lemmas, proofs, and propositions add credibility and depth to the approach.
Experiments are well-designed and span various network configurations, datasets (MNIST, Fashion MNIST, CIFAR-10), and applications (PINNs for solving differential equations). The results consistently show that the proposed method outperforms Xavier initialization, particularly in deeper networks and varying network sizes.

Clarity:
Overall the paper is very well written, with some exceptions mentioned in the weakness section. All the sections in the paper are laid out clearly. The notations are consistent across the paper.

weaknesses:
Significance:
The paper compares their proposed weight initialization with Xavier weight initialization for FCNN with tanh activation. Xavier is known to show diminishing gradients and activations problem for deeper networks, but this is solved by using layer normalization. I am therefore considering this work not significant because the problem that the authors are trying to solve does not exist for Xavier + Layer norm and the authors did not do any comparative analysis with and without layer norm.

Other major issue:
1. In equation 2, the paper claims that a_i^(k+1) follows normal distribution with unit mean. But when number of neurons in layer l-1 (N_(l-1)) is greater than number of neurons in layer l (N_l), then the mean will be greater than 1. When N_(l-1) = 2 * N_l, the mean will be 2. This is not clearly discussed in the paper. If the mean is > 1, then that leads to tanh always saturated.

questions:
1. Layer norm is added to handle the diminishing activation problem. Any reason why you did not compare the performance of the proposed approach with Xavier weight initialization with Layer norm?

2. In equation 2, the paper claims that a_i^(k+1) follows normal distribution with unit mean. But when number of neurons in layer l-1 (N_(l-1)) is greater than number of neurons in layer l (N_l), then the mean can be greater than 1. When N_(l-1) = 2 * N_l, the mean will be 2. This is not clearly discussed in the paper. If the mean is > 1, then that can lead to tanh saturation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729843842556

thread_type:
official_review

<thread_object>
creator: reviewer_kh6z

<content>
summary:
Weight initialisation is an old topic, and most studies have verified that weight initialization methods can improve the performance of neural networks. However, because that neural network’s depth increasing rapidly,  most neural networks, especially Feedforward Neural Networks (FFNNs) should face the gradient vanishment problem.  In this article, authors proposed one weight initialisation method for FFNNs and Physics-Informed Neural Networks (PINNs). Based on an analysis of the fixed points of the function tanh(ax), this method determines values of $a$ that prevent the saturation of activations during the training progress. In terms of robustness, this method presents a stronger and more efficient performance. In the experiment, verified on MNIST, Fashion MNIST, and CIFAR10 datasets, this method also shows acceptable results compared to the Xavier.

strengths:
1. this article proposes one novel weight initialization method for the FFNN and PINN.
2. the authors prove that the activation values cannot vanish when increasing the depth of the neural network by using a fixed-point analysis.

weaknesses:
1. The universality of this method should be improved. For example, more experiments on the feasibility of other neural networks except FNN and PINN.
2. The details of hyperparameters should be mentioned, such as what is the $Threshold$ of FFNN, which training strategy (supervised or unsupervised) of FFNN is used?   
3. FFNN is specifically designed to visualize the trained features, which also should be discussed.
4. The consistency of results should be guaranteed. For example, in Figure 3 (c) and (d), there is a different trend (Xavier tends to equal 
the proposed method), which also should be discussed. In case that after 20 epochs, the performance would be totally different to the presented result.

questions:
The novelty of this work is strong, and the topic sounds interesting. However, the writing and structure should be revised again. There are some questions that the authors should be concerned about. 

1. In Eq. 2, $\sigma_{z}$ is set to $\alpha/\sqrt{N^{l}-1}$ and $\alpha = 0.085$. From Figure 2, we can find that the optimal value of $\alpha$ is $0.085$. Is there any theoretical reason why $\alpha$ should be $0.085$. Or should we manually try the value accordingly?

2. Additionally, please scribe what is $\alpha$. There is no definition of $\alpha$.

3. In Figure 3 (c) and (d), the proposed method seems to decrease after 6 epochs. Although the accuracy curve can rapidly reach the peak (faster than Xavier), the robustness of this method also should be discussed. For example, the initialization method can first provide prior knowledge to neural networks, but if it can keep the stability of training or not. Or is it the reason for the high $\alpha$?

4. In Appendix A.1, the authors discussed different conditions. When $x = 0$, whether the vanishment problem will abscond. Please highlight the strategy on how this method can process it.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730191170857

thread_type:
official_review

<thread_object>
creator: reviewer_6yov

<content>
summary:
This work first provides a theoretical analysis of weight initialization when exclusively using tanh
as an activation function. Providing reasons for the clustering behavior when initializing networks.
Based on the developed theory, an initialization scheme is proposed, and the hyperparameter σz
is empirically determined. There are two types of experiments mainly comparing the proposed
initialization with Xavier. The first type of experiment concerns the classification accuracy in the
early training phase. Showing there is an improvement in accuracy across different data sets and
configurations. The second type of experiment concerns solving PDE with PINNS, showing that the
proposed method has a good performance.

strengths:
Unlike optimization problems with theoretical guarantees on fixed points, weight initialization
is an important task in deep learning. An initialization scheme with theoretical backing can
have a long-lasting impact, even just for a sub-field of deep learning.

Experiments do show significant improvement when using PINNs to solve PDE.

weaknesses:
The impact depended on exclusively using tanh as an activation function is fundamentally
beneficial in PINNs. As the current state of the paper, there is not enough support for this.

Given that the experiments are not too computationally intensive and the experiment section
only considers a few data sets or PDEs, the demonstrated improvement may not be general.

questions:
In sections 4.1 and 4.2, the experiment trains for 20 or 40 epochs. Do networks converge to their best accuracy? What is the difference in accuracy when training for more epochs?

The experiments in section 4 are not too computationally intensive, is it possible to include more
data sets or PDE can show the improvements are general?

Can the code used in the experiment can be provided to improve reproducibility?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730213072568

thread_type:
official_review

<thread_object>
creator: reviewer_79su

<content>
summary:
The paper introduces a novel weight initialization technique specifically designed for neural networks using the tanh activation function. This technique is evaluated against the well-known Xavier initialization method using benchmark datasets. The experimental results demonstrate that the proposed initialization method enhances the convergence speed of Physics-Informed Neural Networks (PINNs) utilizing the tanh function, showing greater robustness to variations in network size. The findings indicate that the new initialization technique outperforms Xavier initialization in solving various problems related to Partial Differential Equations (PDEs).

strengths:
The strength of this paper lies in the development of a novel weight initialization technique that facilitates faster convergence and enhances performance in physics-informed neural networks (PINNs) utilizing the tanh activation function.

weaknesses:
1. The comparison of the proposed weight initialization technique solely with Xavier is insufficient; it should also be experimentally evaluated against other state-of-the-art weight initialization methods.
2. Using tanh activation function in entire neural network is not good practice that it has the drawback of the vanishing gradients for the very high and very low values of x.
3. The formulation is more complex than standard methods, which could complicate implementation as shown in Equation (1). 
4. The optimal value of 𝛼 can be highly context-dependent, varying across different architectures, datasets, and tasks, which makes it less universally applicable. Additionally, the choice of 𝛼 can interact with other hyperparameters, such as learning rate and batch size, complicating the overall tuning process during backpropagation, as described in Equation (2).
5. In Section 4.1, the evaluation process utilizes three datasets—MNIST, FMNIST, and CIFAR-10—employing the tanh activation function in every layer. As shown in Table 2, as the number of hidden layers increases, loss gradually increases, which is indicative of overfitting. It would be more effective to use the proposed weight initialization in conjunction with state-of-the-art architectures for training deep neural networks.

questions:
See above in weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732417796994

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revisions

comment:
Dear Reviewers and AC,

We would like to thank all the reviewers for taking the time to review our work and for providing valuable feedback. 
We appreciate the recognition from reviewers of clear and good presentation **(Reviewers 79Su, 6yoV, kh6z, YCAx)**, improved performance across various tasks **(Reviewers 79Su, 6yoV, YCAx)**, the theoretical foundation of the proposed method **(Reviewers 6yoV, kh6z, YCAx)**, and the novelty of our method **(Reviewers 79Su, kh6z, YCAx)**.
***
The latest revision of our paper has been uploaded, addressing all comments and queries raised by the reviewers. Edits in the PDF have been highlighted in red. Below, we provide a summary of the changes made to our work.

# Writing 
**Improvements** Thanks for the suggestions of Reviewers YCAx, and kh6z
- We have made **minor edits** for clarity and space efficiency.
- We have added **additional explanations in Section 3.1** for further clarification. 
- We have added **inline headings** to each paragraph for improved organization.

# Experiments 
**Normalization Methods** Thanks for the suggestions of Reviewer YCAx
- We have added experiments comparing **Xavier with normalization** and the proposed method on classification datasets in **Figure 4 and Table 3**.
- We have added an experiment comparing **Xavier with normalization** and the proposed method on PDEs in **Table 4**.

**Datasets Efficiency** Thanks for the suggestions of Reviewer 79Su.
- We have added experiments verifying **data efficiency** on classification datasets in **Table 3**.
- We have added experiments verifying **data efficiency** on PDEs in **Figure 6, 7 and 16**.

**Datasets and PDEs**  Thanks for the suggestions of Reviewer 6yoV, kh6z, and 79Su.
- We have added **CIFAR-100 data** to the existing experiments in **Tables 1 and 2**.
- We have added **the diffusion equation and Poisson equation** to the existing experiments in **Table 4**.
- We have added experiments on the **activation value distribution** in deeper layers in **Figures 3 and 8**. 

**Supplementary Experiments**  Thanks for the suggestions of Reviewers 79Su, YCAx, 6yoV, and kh6z.
- We have added experiments on classification datasets between the proposed method for tanh and **He/orthogonal initialization for ReLU in Figure 13**.  
- We have added an experiment on the **autoencoder in Figure 13 (b) and (c)**. 
- We have added a performance comparison experiment of **PINNs based on various activation functions in Figure 10**. 
- We have added an experiment on **PINNs with Swish** activation function in **Figure 18**.

We have included all experimental results in our revised paper.

Best regards,

Authors of submission 13420
</content>
</thread_object>

</thread>
<thread>
creation date:
1732541866341

thread_type:
official_comment

<thread_object>
creator: area_chair_w4e6

<content>
comment:
Dear reviewers,

A reminder that **November, 26** is the last day to interact with the authors, before the private discussion with the area chairs. At the very least, please acknowledge having read the rebuttal (if present). If the rebuttal was satisfying, please improve your score accordingly. Finally, if you have concerns that might be solved in time, this is the last chance before moving on to the next phase.

Thanks,
The AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a novel initialization method for tanh neural networks, based on a fixed-point analysis of the layer.

The reviews are mixed, ranging from marginal rejection to strong acceptance. The reviewers had several technical questions (see below for a detail), but these were addressed in the rebuttal. They were also concerned about some missing baselines (e.g., other initialization techniques, layer normalization), which the authors added in the rebuttal. The only remaining concern is the limited scope of the paper, which focuses on a narrow type of neural network.

Overall, the paper is technically correct; the authors provided a very significant rebuttal which was appreciated by all reviewers who interacted during the rebuttal phase. In addition, they argue correctly that tanh models are still important in several sub-fields, such as PINNs. As a result, I lean towards acceptance.