How to scale weights during training?

username: michaelklachko
cooked: <p>I’d like to train a convnet where each layer weights are divided by the maximum weight in that layer, at the start of every forward pass. So the range of the weights would always be [-1, 1].</p>
<p>I tried doing it like this:</p>
<pre><code class="lang-auto">
class TestConv2d(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel_size=5, bias=False):
        super(TestConv2d, self).__init__(in_channels, out_channels, kernel_size, bias=bias)

    def forward(self, inputs):
        return F.conv2d(inputs, self.weight / torch.max(torch.abs(self.weight)), self.bias)


class TestLinear(nn.Linear):
    def __init__(self, in_features, out_features, bias=False):
        super(TestLinear, self).__init__(in_features, out_features, bias=bias)

    def forward(self, inputs):
        return F.linear(inputs, self.weight / torch.max(torch.abs(self.weight)), self.bias)


class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = TestConv2d(3, 32, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = TestConv2d(32, 64, 5)
        self.flatten = nn.Flatten(start_dim=1)
        self.fc1 = TestLinear(64 * 5 * 5, 390)
        self.fc2 = TestLinear(256, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.flatten(x)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
</code></pre>
<p>The rest of the code is a standard training loop, and it works as expected on CIFAR-10 without weights scaling.</p>
<p>With weight scaling, however, there is some weird behavior. I tried training it with SGD and AdamW, and AdamW works much better. I had to reduce the initial learning rate significantly (compared to no scaling scanario), especially with SGD. However, I’m still not able to reach the same accuracy as without weight scaling (~2% drop with AdamW, ~10% drop with SGD).</p>
<p>Questions:</p>
<ol>
<li>What is happening to the weight gradients when I’m modifying weights like this?</li>
<li>What would be a correct way to implement weight scaling?</li>
<li>Why is AdamW works so much better than SGD in this case?</li>
</ol>

username: ekilic
cooked: <p>First of all, it is very interesting to apply layers with layer scaling. I tried to understand how the deep learning model works with it, but I could not figure out it.</p>
<p>I have some opinion about your questions;<br>
1- Computional graph of your model in the following pic. It is shown how appliying gradinets.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7c455d37175ce73a6ac34b9a1720c877cb30215.png" data-download-href="https://discuss.pytorch.org/uploads/default/e7c455d37175ce73a6ac34b9a1720c877cb30215" title="backw"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7c455d37175ce73a6ac34b9a1720c877cb30215_2_225x375.png" alt="backw" data-base62-sha1="x4iWp8uY3jdO3YqTZjnkspUol4F" width="225" height="375" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7c455d37175ce73a6ac34b9a1720c877cb30215_2_225x375.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7c455d37175ce73a6ac34b9a1720c877cb30215_2_337x562.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7c455d37175ce73a6ac34b9a1720c877cb30215.png 2x" data-dominant-color="EAECEB"><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">backw</span><span class="informations">405×673 22.3 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>2- I have no idea.</p>
<p>3- Did you use L2 regularization with SGD? AdamW applies built-in L2 regularization.</p>

username: michaelklachko
cooked: <p>Thank you!</p>
<ol>
<li>
<p>How did you make this diagram?</p>
</li>
<li>
<p>From the graph, it seems like the weights receive two sets of gradients: one normal, and one that goes through max/abs ops, so I’m guessing it gets applied only to the largest weight? Very interesting indeed.</p>
</li>
<li>
<p>I did use L2 regularization for both SGD and AdamW. I tried different values, and it doesn’t seem to be very sensitive to it, but some values are better than others.</p>
</li>
</ol>

username: googlebot
cooked: <p>This reminds me of <a href="https://github.com/zj10/ND-Adam" rel="nofollow noopener">Normalized Direction-preserving Adam</a>. Regarding your approach, you should probably use vector maximums instead of a single scalar divisor.<br>
Also, <a href="https://pytorch.org/docs/stable/nn.html#weight-norm" rel="nofollow noopener">nn.utils.weight_norm</a> may be relevant.</p>

username: michaelklachko
cooked: <p>Thank you! What do you mean “vector maximums”?</p>

username: michaelklachko
cooked: <p>Looking at the computational graph, it seems like the largest weight in each layer might be receiving two values of gradients, one normal, and one coming through the path with max and abs ops. Any idea how to verify that?</p>

username: googlebot
cooked: <aside class="quote no-group" data-username="michaelklachko" data-post="5" data-topic="90796" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/michaelklachko/48/541_2.png" class="avatar"> michaelklachko:</div>
<blockquote>
<p>Thank you! What do you mean “vector maximums”?</p>
</blockquote>
</aside>
<p>torch.max(weights, dim=1,keepdim=True)[0]</p>

username: michaelklachko
cooked: <p>Sorry, can you please explain why this is better?</p>

username: googlebot
cooked: <p>Can you explain why is matrix max norm better that vector (L2?) norms? My understanding is that using vector norms (either rowwise or columnwise) for normalization don’t “tie” features together, so outputs are more expressive/diversified, with balanced variance.</p>

username: michaelklachko
cooked: <p>Are you talking about channel range equalization, i.e. dividing by the channel max instead of the layer max? That’s an interesting idea, but do we actually want to “untie” the features captured by different channels? If a channel is “dead” with weights very close to zero, do we want to rescale it to be as strong as other, useful channels? Simply scaling down all channels by the same values preserves the relative channel importance information. I wonder if batch normalization does some form of channel equalization, but I believe it still allows some channels to go “dead”.</p>
<p>Anyway, I just tried it:</p>
<pre><code class="lang-auto">F.conv2d(inputs, self.weight / torch.abs(self.weight).max(dim=1, keepdim=True)[0].max(dim=2, keepdim=True)[0].max(dim=3, keepdim=True)[0])
F.linear(inputs, self.weight / torch.abs(self.weight).max(dim=1, keepdim=True)[0])
</code></pre>
<p>From what I see so far after adjusting learning rate and weight decay it does not train as well as dividing by the layer max weight (~3% drop) while I was able to reach almost the baseline accuracy with layer max scaling (~0.5% drop). The training with per channel scaling is not as stable as layer wise scaling, and is more sensitive to the hyperparam choices. It’s possible that with more effort I will close the gap between these two scaling methods. Though I still don’t quite understand what happens to gradients, the computational graph does not make much sense to me.</p>

username: ekilic
cooked: <p>I used pytorchviz for diagram. (<a href="https://github.com/szagoruyko/pytorchviz" rel="nofollow noopener">https://github.com/szagoruyko/pytorchviz</a>)</p>

username: googlebot
cooked: <p>All weight normalization variants restrict magnitudes, so there should be some layer that can rescale features freely.</p>
<aside class="quote no-group" data-username="michaelklachko" data-post="10" data-topic="90796">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/michaelklachko/48/541_2.png" class="avatar"> michaelklachko:</div>
<blockquote>
<p>If a channel is “dead” with weights very close to zero, do we want to rescale it to be as strong as other, useful channels? Simply scaling down all channels by the same values preserves the relative channel importance information. I wonder if batch normalization does some form of channel equalization, but I believe it still allows some channels to go “dead”.</p>
</blockquote>
</aside>
<p>It is an intermediate representation, so there is no value in “dead” channels. On the contrary, you get problems with vanishing gradients in multi-layer pipelines. Indeed, batch normalization helps to avoid “dead” channels (affine transformation it does is kinda auxiliary, I think).</p>
<aside class="quote no-group" data-username="michaelklachko" data-post="10" data-topic="90796">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/michaelklachko/48/541_2.png" class="avatar"> michaelklachko:</div>
<blockquote>
<p>Though I still don’t quite understand what happens to gradients, the computational graph does not make much sense to me.</p>
</blockquote>
</aside>
<p>You have a special cell in a weight matrix (located at argmax()) that all channels use to try to adjust outputs. Row that contains this cell thus produces a feature with different gradient magnitude and/or variance. So, SGD suffers, while Adam adjusts to mitigate this.</p>

username: mohit117
cooked: <p>The solution to your original question is very straightforward. The key idea is to use <code>torch.nn.Parameter</code> as shown below.</p>
<pre><code class="lang-auto">class MyCNN(torch.nn.Module):
    def __init__(self):
        super(MyCNN, self).__init__()
        
        conv1 = nn.Conv2d(..specify kernel size, stride, etc..)
        self.conv1_w = torch.nn.Parameter(conv1.weight)
        self.conv1_b = torch.nn.Parameter(conv1.bias)
        
        # have several more like this.
        
    def forward(self, x):
        
        x = torch.nn.functional.conv2d(x, self.conv1_w, self.conv1_b, stride=1, padding=0, dilation=1, groups=1)
        
        self.conv1_w = self.conv1_b/torch.max(self.conv1_w)
        self.conv1_b = self.conv1_b/torch.max(self.conv1_b)
        
        return x
</code></pre>