Systematic Study of Numerical Instability Caused by Extreme Weight Initialization in PyTorch

zhouxiaoyaozzz: <p>Problem Description</p>
<p>I’m investigating how weight initialization affects training stability in neural networks. Two architecturally similar models show drastically different behaviors:</p>
<pre><code class="lang-auto"># Stable version
class SimpleModel(nn.Module):
    def __init__(self, input_dim, hidden_dim=128, output_dim=2):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)  # Default init
        self.layer2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, inputs):
        with torch.amp.autocast(enabled=False, device_type="cuda"):
            x = self.layer1(inputs)
        x = self.layer2(x)
        return x

# Unstable version
class UnstableModel(nn.Module):
    def __init__(self, input_dim, hidden_dim=128, output_dim=2):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.layer2 = nn.Linear(hidden_dim, output_dim)
        
        # Extreme initialization
        nn.init.uniform_(self.layer1.weight, -100, 100)  # Range &gt;&gt; typical
        nn.init.zeros_(self.layer1.bias)
        nn.init.uniform_(self.layer2.weight, -100, 100)
        nn.init.zeros_(self.layer2.bias)

    def forward(self, inputs):
        x = self.layer1(inputs)  # No stability measures
        x = self.layer2(x)
        return x
</code></pre>
<p>When comparing two architecturally identical models with different initializations, we observe outputs differing by 4 orders of magnitude:</p>
<pre><code class="lang-auto">Input = [0.1, 0.2, ..., 1.0]  # 10D features

# Output comparison
Stable_model = [-0.13, -0.05]      # Magnitude: 10^0
Extreme_init_model = [28528, 55630] # Magnitude: 10^4
</code></pre>
<p>Questions:</p>
<ol>
<li>Why does extreme weight initialization (e.g., ∈[-100,100]) cause exponential output growth?</li>
<li>How to identify the layer where instability first occurs?</li>
<li>What does this reveal about model robustness?</li>
</ol>

KFrank: <p>Hi Xiaoyao!</p>
<aside class="quote no-group quote-modified" data-username="zhouxiaoyaozzz" data-post="1" data-topic="218356">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zhouxiaoyaozzz/48/76376_2.png" class="avatar"> zhouxiaoyaozzz:</div>
<blockquote>
<p>I’m investigating how weight initialization affects training stability<br>
…</p>
<pre><code class="lang-auto">        # Extreme initialization
        nn.init.uniform_(self.layer1.weight, -100, 100)  # Range &gt;&gt; typical
        nn.init.zeros_(self.layer1.bias)
        nn.init.uniform_(self.layer2.weight, -100, 100)
        nn.init.zeros_(self.layer2.bias)
</code></pre>
<p>…</p>
<pre><code class="lang-auto"># Output comparison
Stable_model = [-0.13, -0.05]      # Magnitude: 10^0
Extreme_init_model = [28528, 55630] # Magnitude: 10^4
</code></pre>
<p>…</p>
<ol>
<li>Why does extreme weight initialization (e.g., ∈[-100,100]) cause exponential output growth?</li>
</ol>
</blockquote>
</aside>
<p>First comment:</p>
<p>If you pass inputs of order one through a <code>Linear</code> whose weights are of order 10^2, you<br>
would expect to get outputs of order 10^2.  Passing those inputs through two order-10^2<br>
<code>Linear</code>s would naturally produce outputs of order 10^4.</p>
<aside class="quote no-group">
<blockquote>
<ol start="2">
<li>How to identify the layer where instability first occurs?</li>
</ol>
</blockquote>
</aside>
<p>Second comment:</p>
<p>You asked about training stability, but you haven’t told us anything about how (or whether)<br>
you trained your model.</p>
<p>However, gradient-descent optimization routinely becomes unstable when the steps taken<br>
are too large.  The size of a step is the size of the gradient times the learning rate.  If the<br>
resulting step is too large, you jump to a place where the weights are worse and the<br>
gradient is larger.  Then you take a bigger jump to where the weights are even worse and<br>
the gradient is even larger and the training systematically diverges.</p>
<p>When you initialize the weights to large values, the gradients will start out large and you<br>
risk this kind of divergence / instability.  You can compensate for this by starting with a<br>
smaller learning rate, say one that is smaller by a factor of 10^-4.</p>
<p>(Note, if you systematically increase the learning rate with which you train the stable<br>
version of your model, you will be able to make its training similarly unstable.)</p>
<aside class="quote no-group">
<blockquote>
<ol start="3">
<li>What does this reveal about model robustness?</li>
</ol>
</blockquote>
</aside>
<p>This doesn’t really say anything about the robustness of the model itself.  It’s the mismatch<br>
between the learning rate and the perversely-large initial values of the model’s weights that<br>
is causing the unstable training.</p>
<p>(In your case, you should be able to start training your “unstable” model with a small<br>
learning rate and after the weights have evolved to more reasonable values, continue<br>
training with a larger, more typical learning rate.)</p>
<p>Best.</p>
<p>K. Frank</p>