Use of BatchNorm1d seems to prevent some model parameters from being updated

username: Jim_Thompson
cooked: <p>Looking for guidance around the use of <code>BatchNorm1d</code> and how its use affect parameter updating.</p>
<p>Are there are any guidelines on how to use <code>BatchNorm1d</code>?</p>
<p>The reason I’m asking is that it appears when I use <code>BatchNorm1d</code>, some model parameters are not updated during the backward pass.</p>
<p>Here is a minimal example that illustrates the issue:</p>
<pre><code class="lang-auto">import torch

INPUT_SIZE = 2
OUTPUT_SIZE = 2
BATCH_SIZE = 4
NUM_STEPS = 2


class SubModule2WithBN(torch.nn.Module):
    def __init__(self, size):
        super().__init__()
        self.fc = torch.nn.Linear(size, size)
        self.bn = torch.nn.BatchNorm1d(size)

    def forward(self, inputs):
        hidden = self.fc(inputs)
        return self.bn(hidden)


class SubModule2WithoutBN(torch.nn.Module):
    def __init__(self, size):
        super().__init__()
        self.fc = torch.nn.Linear(size, size)

    def forward(self, inputs):
        return self.fc(inputs)


class MyModel(torch.nn.Module):

    def __init__(self, input_size, output_size):
        super().__init__()

        self.input_size = input_size
        self.output_size = output_size

        self.fc1 = torch.nn.Linear(input_size, output_size)
        self.fc2 = torch.nn.Linear(output_size, output_size)

        self.submodule2 = torch.nn.ModuleList()

        for _ in range(NUM_STEPS):
            # Comment out one or the other of the following statements to run this minimal example
            self.submodule2.append(SubModule2WithBN(output_size))   # this will result in some parameters not updated
            # self.submodule2.append(SubModule2WithoutBN(output_size))  # this will result in all parameters updated

    def forward(self, inputs):
        hidden = self.fc1(inputs)
        for step in range(NUM_STEPS):
            hidden = self.submodule2[step](hidden)

        return self.fc2(hidden)


# generate synthetic tensor to feed into model
torch.random.manual_seed(1919)
input_tensor = torch.randn([BATCH_SIZE, INPUT_SIZE], dtype=torch.float32)

# create model instance
model = MyModel(INPUT_SIZE, OUTPUT_SIZE)
model.train(True)

# run synthetic input tensor through the model
model_output = model(input_tensor)

#
# Purpose of the following code is to test if all parameters were updated
# not to achieve an optimal solution.
#
# setup
loss_function = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
optimizer.zero_grad()

# generate psuedo target tensor used for calculating loss
target_tensor = torch.randn(model_output.shape, dtype=model_output.dtype)

# capture model parameters before doing parameter update pass
before = [(x[0], x[1].clone()) for x in model.named_parameters()]

# do update of model parameters
loss = loss_function(model_output, target_tensor)
loss.backward()
optimizer.step()

# capture model parameters after pass for updating parameters
after = [(x[0], x[1].clone()) for x in model.named_parameters()]

# confirm that at all parameters were updated
parameters_updated = []
parameter_update_results = []
for b, a in zip(before, after):
    parameter_updated = (a[1] != b[1]).any()
    parameters_updated.append(parameter_updated)
    parameter_update_results.append(
        f'\nParameter {a[0]} updated: {parameter_updated}{" &lt;&lt;&lt;&lt;&lt;&lt;&lt;" if not parameter_updated else ""}\n'
        f'before model forward() pass (requires grad:{b[1].requires_grad}): \n\t{b[1]}\n'
        f'after model forward() pass (requires grad:{a[1].requires_grad}): \n\t{a[1]}\n'
        f'|before - after|: \n\t{torch.abs(b[1] - a[1])}\n'
    )
print(
    f'\nAll parameters updated: {all(parameters_updated)} {" &lt;&lt;&lt;&lt;&lt;&lt;&lt;" if not all(parameters_updated) else ""}'
    f'{"".join(parameter_update_results)}'
)
</code></pre>
<p>In the above example code, if I use the class <code>SubModule2WithBN</code>, which contains <code>Batchnorm1d</code>, I see output like this:</p>
<pre><code class="lang-auto">All parameters updated: False  &lt;&lt;&lt;&lt;&lt;&lt;&lt;
Parameter fc1.weight updated: True
before model forward() pass (requires grad:True): 
	tensor([[-0.5268,  0.5782],
        [-0.4338, -0.2834]], grad_fn=&lt;CloneBackward0&gt;)
after model forward() pass (requires grad:True): 
	tensor([[-0.5336,  0.5678],
        [-0.4505, -0.2665]], grad_fn=&lt;CloneBackward0&gt;)
|before - after|: 
	tensor([[0.0068, 0.0104],
        [0.0167, 0.0169]], grad_fn=&lt;AbsBackward0&gt;)

Parameter fc1.bias updated: False &lt;&lt;&lt;&lt;&lt;&lt;&lt;
before model forward() pass (requires grad:True): 
	tensor([0.5317, 0.6025], grad_fn=&lt;CloneBackward0&gt;)
after model forward() pass (requires grad:True): 
	tensor([0.5317, 0.6025], grad_fn=&lt;CloneBackward0&gt;)
|before - after|: 
	tensor([0., 0.], grad_fn=&lt;AbsBackward0&gt;)
	
&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt; deleted lines that showed parameters were updated &gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;

Parameter submodule2.1.fc.bias updated: False &lt;&lt;&lt;&lt;&lt;&lt;&lt;
before model forward() pass (requires grad:True): 
	tensor([-0.5061,  0.0665], grad_fn=&lt;CloneBackward0&gt;)
after model forward() pass (requires grad:True): 
	tensor([-0.5061,  0.0665], grad_fn=&lt;CloneBackward0&gt;)
|before - after|: 
	tensor([0., 0.], grad_fn=&lt;AbsBackward0&gt;)

Parameter submodule2.1.bn.weight updated: True
before model forward() pass (requires grad:True): 
	tensor([1., 1.], grad_fn=&lt;CloneBackward0&gt;)
after model forward() pass (requires grad:True): 
	tensor([0.9230, 0.9826], grad_fn=&lt;CloneBackward0&gt;)
|before - after|: 
	tensor([0.0770, 0.0174], grad_fn=&lt;AbsBackward0&gt;)

Parameter submodule2.1.bn.bias updated: True
before model forward() pass (requires grad:True): 
	tensor([0., 0.], grad_fn=&lt;CloneBackward0&gt;)
after model forward() pass (requires grad:True): 
	tensor([-0.0509, -0.0126], grad_fn=&lt;CloneBackward0&gt;)
|before - after|: 
	tensor([0.0509, 0.0126], grad_fn=&lt;AbsBackward0&gt;)


Process finished with exit code 0

</code></pre>
<p>The above show two parameters not updated.</p>
<p>OTOH, if I use this class <code>SubModule2WithoutBN</code>, which does not have <code>Batchnorm1d</code>, then all parameters are updated:</p>
<pre><code class="lang-auto">All parameters updated: True 
Parameter fc1.weight updated: True
before model forward() pass (requires grad:True): 
	tensor([[-0.5268,  0.5782],
        [-0.4338, -0.2834]], grad_fn=&lt;CloneBackward0&gt;)
after model forward() pass (requires grad:True): 
	tensor([[-0.5212,  0.5742],
        [-0.4376, -0.2759]], grad_fn=&lt;CloneBackward0&gt;)
|before - after|: 
	tensor([[0.0057, 0.0040],
        [0.0038, 0.0075]], grad_fn=&lt;AbsBackward0&gt;)

Parameter fc1.bias updated: True               # &lt;= this is now updated
before model forward() pass (requires grad:True): 
	tensor([0.5317, 0.6025], grad_fn=&lt;CloneBackward0&gt;)
after model forward() pass (requires grad:True): 
	tensor([0.5322, 0.5955], grad_fn=&lt;CloneBackward0&gt;)
|before - after|: 
	tensor([0.0005, 0.0070], grad_fn=&lt;AbsBackward0&gt;)

&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt; deleted lines that showed parameters were updated &gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;

Parameter submodule2.1.fc.weight updated: True
before model forward() pass (requires grad:True): 
	tensor([[-0.3351, -0.6894],
        [ 0.2473, -0.2471]], grad_fn=&lt;CloneBackward0&gt;)
after model forward() pass (requires grad:True): 
	tensor([[-0.3261, -0.6792],
        [ 0.2027, -0.2713]], grad_fn=&lt;CloneBackward0&gt;)
|before - after|: 
	tensor([[0.0090, 0.0102],
        [0.0445, 0.0242]], grad_fn=&lt;AbsBackward0&gt;)

Parameter submodule2.1.fc.bias updated: True           #&lt;= this is now updated
before model forward() pass (requires grad:True): 
	tensor([-0.5061,  0.0665], grad_fn=&lt;CloneBackward0&gt;)
after model forward() pass (requires grad:True): 
	tensor([-0.4968,  0.0194], grad_fn=&lt;CloneBackward0&gt;)
|before - after|: 
	tensor([0.0093, 0.0471], grad_fn=&lt;AbsBackward0&gt;)


Process finished with exit code 0
</code></pre>
<p>Some specific questions:</p>
<ul>
<li>Am I using <code>BatchNorm1d</code> correctly?</li>
<li>Is this a reasonable way to test for updating of parameters?</li>
</ul>
<pre><code class="lang-auto">#
# Purpose of the following code is to test if all parameters were updated
# not to achieve an optimal solution.
#
# setup
loss_function = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
optimizer.zero_grad()

# generate psuedo target tensor used for calculating loss
target_tensor = torch.randn(model_output.shape, dtype=model_output.dtype)

# capture model parameters before doing parameter update pass
before = [(x[0], x[1].clone()) for x in model.named_parameters()]

# do update of model parameters
loss = loss_function(model_output, target_tensor)
loss.backward()
optimizer.step()

# capture model parameters after pass for updating parameters
after = [(x[0], x[1].clone()) for x in model.named_parameters()]

</code></pre>
<p>Any guidance will be appreciated.</p>

username: ptrblck
cooked: <p>Thanks for the great code and the interesting question!<br>
Based on your output (and my runs) it seems that the bias of the preceding linear layer is not updated.<br>
While the check claims so, you could check the gradients of the <code>bias</code> parameter and would see that these gradients are indeed really small:</p>
<pre data-code-wrap="python"><code class="lang-python">model.fc1.bias.grad
&gt; tensor([ 5.9605e-08, -1.7881e-07])
</code></pre>
<p>and thus the bias value is not changed.</p>
<p>Now, this might be surprising, but let’s think about the first operation of a batchnorm layer:</p>
<pre data-code-wrap="python"><code class="lang-python">out = (x - mean) / stddev * weight + bias
</code></pre>
<p>As you can see, the mean of the incoming batch would be subtracted. If the previous layer has added a bias to the activation, it would be directly subtracted again.<br>
Wouldn’t this also mean that this parameter has (almost) zero influence on the loss calculation (up to the limited numerical precision limit)?<br>
The <a href="https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#disable-bias-for-convolutions-directly-followed-by-a-batch-norm">performance guide</a> also mentions:</p>
<blockquote>
<p>If a <code>nn.Conv2d</code> layer is directly followed by a <code>nn.BatchNorm2d</code> layer, then the bias in the convolution is not needed, instead use <code>nn.Conv2d(..., bias=False, ....)</code>. Bias is not needed because in the first step <code>BatchNorm</code> subtracts the mean, which effectively cancels out the effect of bias.<br>
This is also applicable to 1d and 3d convolutions as long as <code>BatchNorm</code> (or other normalization layer) normalizes on the same dimension as convolution’s bias.</p>
</blockquote>
<p>However, this targets mainly the performance and would thus avoid the unnecessary <code>add</code> kernel.</p>
<p>CC <a class="mention" href="/u/tom">@tom</a> and <a class="mention" href="/u/kfrank">@KFrank</a> to correct me here in case my explanation is wrong. I’m also sure both can add interesting explanations from a more mathematical point of view. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>

username: tom
cooked: <p>As you might expect, <a class="mention" href="/u/ptrblck">@ptrblck</a> 's explanation totally nails it, I’d probably strike the “almost” in almost zero influence. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"> The gradients should be 0 here as the result does not depend on the parameter. In other words, use bias=False for the linear/conv preceding batch norm.</p>
<p>If it is any consolation, it is <span class="hashtag">#5</span> on Andrej Karpathy’s <a href="https://twitter.com/karpathy/status/1013245864570073090" rel="noopener nofollow ugc">somewhat famous list of common NN training mistakes</a>, so you are in good company.</p>
<p>Best regards</p>
<p>Thomas</p>

username: Jim_Thompson
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>  <a class="mention" href="/u/tom">@tom</a>   thank you very much for a quick and clear explanation.</p>
<p>When I use <code>bias=False</code> in the the preceding <code>Linear</code> layer, I see the indicator that all parameters have been updated.</p>
<p>I believe your response explains what I saw in a subsequent test I ran after posting my question.  In this later test, instead of doing only one backward pass with parameter updating, I ran multiple iterations.  In this test, I eventually saw the bias parameter updated after multiple iterations.  Sometimes it took up to 200 iterations for an update to be detected.</p>
<p>In other words, it took many iterations for the the effect of the very small <code>bias</code> gradients to “accumulate” enough to have an effect on parameter updating.</p>
<p>Again, thank you for your insights and the link to PyTorch performance guide.</p>