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”.
Anyway, I just tried it:
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])

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.