<p>Looks great, as you have already followed most of the solutions to resolve gradient exploding problem. Below is the list of all solutions you can try</p>
<p><strong>Solutions to avoid Gradient Exploding problem</strong></p>
<ol>
<li><p><em>Appropriate Weight initialization:</em> utilise appropriate weight Initialization based on the activation function used.</p>
<div class="s-table-container">
<table class="s-table">
<thead>
<tr>
<th>Initialization</th>
<th>Activation Function</th>
</tr>
</thead>
<tbody>
<tr>
<td>He</td>
<td>ReLU &amp; variants</td>
</tr>
<tr>
<td>LeCun</td>
<td>SELU</td>
</tr>
<tr>
<td>Glorot</td>
<td>Softmax, Logistic, None, Tanh</td>
</tr>
</tbody>
</table>
</div></li>
<li><p><em>Redesigning your Neural network:</em> use fewer layers in neural network and/or use smaller batch size</p>
</li>
<li><p><em>Choosing Non Saturation activation function</em>: choose the right activation function with reduced learning rates</p>
<ul>
<li>ReLU</li>
<li>Leaky ReLU</li>
<li>randomized leaky ReLU (RReLU)</li>
<li>parametric leaky ReLU (PReLU)</li>
<li>exponential linear unit (ELU)</li>
</ul>
</li>
<li><p><em>Batch Normalisation:</em> Ideally using batch normalisation before/after each layer, based on what works best for your dataset.</p>
<ul>
<li><p>after each layer <a href="https://arxiv.org/abs/1502.03167" rel="nofollow noreferrer">Paper reference</a></p>
<pre><code>model = keras.models.Sequential([
                     keras.layers.Flatten(input_shape=[28, 28]),
                     keras.layers.BatchNormalization(),
                     keras.layers.Dense(300, activation=&quot;elu&quot;, 
                     kernel_initializer=&quot;he_normal&quot;),
                     keras.layers.BatchNormalization(),
                     keras.layers.Dense(100, activation=&quot;elu&quot;, 
                     kernel_initializer=&quot;he_normal&quot;),
                     keras.layers.BatchNormalization(),
                     keras.layers.Dense(10, activation=&quot;softmax&quot;)
         ])
</code></pre>
</li>
<li><p>before each layer</p>
<pre><code> model = keras.models.Sequential([
                     keras.layers.Flatten(input_shape=[28, 28]),
                     keras.layers.BatchNormalization(),
                     keras.layers.Dense(300, kernel_initializer=&quot;he_normal&quot;, use_bias=False),
                     keras.layers.BatchNormalization(),
                     keras.layers.Activation(&quot;elu&quot;),
                     keras.layers.Dense(100, kernel_initializer=&quot;he_normal&quot;, use_bias=False),
                     keras.layers.Activation(&quot;elu&quot;),
                     keras.layers.BatchNormalization(),
                     keras.layers.Dense(10, activation=&quot;softmax&quot;)
         ])
</code></pre>
</li>
</ul>
</li>
<li><p><em>Gradient Clipping</em> : Good default values are clipnorm=1.0 and clipvalue=0.5</p>
</li>
<li><p><em>Ensure right optimizer is utilised</em>: Since you have utilised Adam optimizer, check if other optimizer works best for your case. Refer <a href="https://keras.io/api/optimizers/" rel="nofollow noreferrer">this documentation</a> for info on the available optimizers [SGD, RMSprop, Adam, Adadelta, Adagrad, Adamax, Nadam, Ftrl]</p>
</li>
<li><p><em>Truncated Backpropagation through time</em>: often works for RNNS refer this <a href="https://machinelearningmastery.com/gentle-introduction-backpropagation-time/" rel="nofollow noreferrer">documentation</a></p>
</li>
<li><p><em>Use LSTM</em>(solution for RNN)</p>
</li>
<li><p><em>Use weight regularizers on layers</em>: set <code>kernel_regularizer</code> to L1 or L2. <a href="https://keras.io/api/layers/regularizers/" rel="nofollow noreferrer">Weight regularizer document reference</a></p>
</li>
</ol>
<p>For more information refer to chapter 11 in <em>Hands on Machine learning with scikit-learn, keras and tensorflow</em> book by <em>Aurélien</em></p>
Gradient exploding problem in a graph neural network
<p>I have a gradient exploding problem which I couldn't solve after trying for several days. I implemented a custom message passing graph neural network in TensorFlow which is used to predict a continuous value from graph data. Each graph is associated with one target value. Each node of a graph is represented by a node attribute vector, and the edges between nodes are represented by an edge attribute vector.</p>
<p>Within a message passing layer, node attributes are updated in a certain way (e.g., by aggregating other node/edge attributes), and these updated node attributes are returned.</p>
<p>Now, I managed to figure out where the gradient problem occurs in my code. I have the below snippet.</p>
<pre><code>to_concat = [neighbors_mean, e]
z = K.concatenate(to_concat, axis=-1)
output = self.Net(z)
</code></pre>
<p>Here, <code>neighbors_mean</code> is the element-wise mean between two node attributes <code>vi</code>, <code>vj</code> that form the edge having an edge attribute <code>e</code>. <code>Net</code> is a single layer feed-forward network. With this, the training loss suddenly jumps to NaN after about 30 epochs with a batch size of 32. If the batch size is 128, still the gradients explode after about 200 epochs.</p>
<p>I found that, in this case, the gradients explode because of the edge attribute <code>e</code>. If I didn't concatenate <code>neighbors_mean</code> with <code>e</code> and just used the below code, there would be no gradient explosion.</p>
<pre><code>output = self.Net(neighbors_mean)
</code></pre>
<p>Also I can avoid gradient explosion by sending <code>e</code> through a sigmoid function as follows. But this degrades the performance (final MAE), because the values in <code>e</code> are mapped to 0-1 range non-linearly. Note that <a href="https://en.wikipedia.org/wiki/Rectifier_(neural_networks)" rel="nofollow noreferrer">Rectified Linear Unit</a> (ReLU) instead of sigmoid didn't work.</p>
<pre><code>to_concat = [neighbors_mean, tf.math.sigmoid(e)]
z = K.concatenate(to_concat, axis=-1)
output = self.Net(z)
</code></pre>
<p>Just to mention that <code>e</code> carries a single value relating to the distance between the two corresponding nodes and this distance is always in the range 0.5-4. There are no large values or NaNs in <code>e</code>.</p>
<p>I have a custom loss function to train this model, but I found that this is not a problem with loss (other losses also led to the same problem). Below is my custom loss function. Note that although this is a single output regression network, the final layer of my NN has two neurons, relating to the mean and log(sigma) of the prediction.</p>
<pre><code>def robust_loss(y_true, y_pred):
  &quot;&quot;&quot;
  Computes the robust loss between labels and predictions.
  &quot;&quot;&quot;
  mean, sigma = tf.split(y_pred, 2, axis=-1)
  # tried limiting 'sigma' with  sigma = tf.clip_by_value(sigma,-4,1.0) but the gradients still explode
  loss =  np.sqrt(2.0) * K.abs(mean - y_true) * K.exp(-sigma)  + sigma
  return K.mean(loss)
</code></pre>
<p>I basically tried everything suggested online to avoid gradient explosion.</p>
<ol>
<li>Applied gradient clipping - with <code>Adam(lr, clipnorm=1, clipvalue=5)</code> and also with <code>tf.clip_by_global_norm(gradients, 1.0)</code></li>
<li>My target variables are always scaled</li>
<li>Weights are initialized with <code>glorot_uniform</code> distribution</li>
<li>Applied regularisation to weights</li>
<li>Tried larger batch sizes (till 256, although delayed gradient explosion happens at some point)</li>
<li>Tried with reduced learning rate</li>
</ol>
<p>What am I missing here? I definitely know it has something to do with concatenating <code>e</code>. But given that 0.5&lt;e&lt;4, why do the gradients explode in this case? This feature <code>e</code> is important to me. What else can I do to avoid numerical overflow in my model?</p>