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].
I tried doing it like this:

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

The rest of the code is a standard training loop, and it works as expected on CIFAR-10 without weights scaling.
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).
Questions:

What is happening to the weight gradients when I’m modifying weights like this?
What would be a correct way to implement weight scaling?
Why is AdamW works so much better than SGD in this case?