Can't figure out why my internal layers in GNN have near-zero gradient

Coprosmo: <p>Hi guys, I’m training a graph neural network of the form (Fully connected layers → Message-passing steps → Fully connected layers), and have found that while most of the fully connected layers are learning, the message-passing layers have extremely low gradient.</p>
<p>I thought this might be a case of the vanishing gradient problem, though I’m not sure since the fully-connected layers before the message-passing layers are learning and have non-zero gradient.</p>
<p>I discovered the issue by looking at the Weights&amp;Biases parameter/gradient visualization, and found that the gradients on the message-passing layers were all close to zero. As a result, the parameters barely shift from their initial values.</p>
<p>Note in the below images that the first set’s gradient values are an order of magnitude lower than the second.</p>
<p>Gradient values in one of the linear layers in the first message-passing layer:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/0/00c548262c5ed22b90122e4555a4e860adf8027f.png" data-download-href="https://discuss.pytorch.org/uploads/default/00c548262c5ed22b90122e4555a4e860adf8027f" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/0/00c548262c5ed22b90122e4555a4e860adf8027f.png" alt="image" data-base62-sha1="6OFOiRbxosFB1m8ZqxnewEAmnR" width="690" height="252" data-dominant-color="F1F4F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">791×290 16.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Resulting parameter distribution:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/c/cc35fd5d2a6f07707c8048204f5e1e107cf598a1.png" data-download-href="https://discuss.pytorch.org/uploads/default/cc35fd5d2a6f07707c8048204f5e1e107cf598a1" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/c/cc35fd5d2a6f07707c8048204f5e1e107cf598a1.png" alt="image" data-base62-sha1="t8x4JAExaXzrCB7xPEjtsidD8Yx" width="690" height="257" data-dominant-color="8DBEE4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">790×295 26.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Gradient values in first fully-connected layer of network:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/5/9594a6d36936c02282636ffd66f1811ebaa26ff5.png" data-download-href="https://discuss.pytorch.org/uploads/default/9594a6d36936c02282636ffd66f1811ebaa26ff5" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/5/9594a6d36936c02282636ffd66f1811ebaa26ff5.png" alt="image" data-base62-sha1="llfCVkkNU5BfVez0XsxVS8HBUDr" width="690" height="257" data-dominant-color="F0F5F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">781×292 15.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Resulting parameter distribution:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/0/902a0ac46872c1b01760100b597d2cc4a239e7a7.png" data-download-href="https://discuss.pytorch.org/uploads/default/902a0ac46872c1b01760100b597d2cc4a239e7a7" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/0/902a0ac46872c1b01760100b597d2cc4a239e7a7.png" alt="image" data-base62-sha1="kzkPuduhjJlFX2zDSycGDhWD7z9" width="690" height="249" data-dominant-color="D6E5F2"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">790×286 17.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>However, some layers in the message-passing steps do seem to learn. Here’s one which has higher gradient values:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/5/25ad8c7c7c8765c8929bab6d579b6dd9500ba8e5.png" data-download-href="https://discuss.pytorch.org/uploads/default/25ad8c7c7c8765c8929bab6d579b6dd9500ba8e5" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/5/25ad8c7c7c8765c8929bab6d579b6dd9500ba8e5.png" alt="image" data-base62-sha1="5njuyRGePb71W731b5pKLlrnJWd" width="690" height="248" data-dominant-color="EFF4F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">794×286 15.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Resulting parameter distribution:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/4/e4c361ff96e995c5fef56c83e7af98017fb23d05.png" data-download-href="https://discuss.pytorch.org/uploads/default/e4c361ff96e995c5fef56c83e7af98017fb23d05" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/4/e4c361ff96e995c5fef56c83e7af98017fb23d05.png" alt="image" data-base62-sha1="wDJt4OreGIRQ89q4AVfz4iYxVSl" width="690" height="253" data-dominant-color="C0D9EC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">786×289 24.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Does anyone have a clue as to what might be going on here? Why would the gradient values at some of the layers of the message-passing nn be so low?</p>
<p>I am training using SGD with LR=0.01, momentum=0.8, dropout=0.5 after all message-passing fully-connected layers. To code the graph neural-network, I am using PyTorch Geometric.</p>