Removing spectral norm changes the output values

username: GuillaumeTong
cooked: <p>In my current work, I need to use spectral norms (SN), and I want to export my model to onnx, unfortunately, SN messes with the onnx export.<br>
When trying to remove the SN, export to onnx and put SN back on, I realized that removing the spectral norm and putting in back in changes the weights, and their effects on the output.</p>
<p>Let me illustrate this with the following sample code:</p>
<pre><code class="lang-auto">import torch


t0 = torch.randn(2000, 3, 4, 4)
c1 = torch.nn.Conv2d(3, 1, 3)
w1 = c1.weight.data.clone()
o1 = c1(t0)


c2 = torch.nn.utils.spectral_norm(c1)
w2 = c2.weight.data.clone()
o2 = c2(t0)


c3 = torch.nn.utils.remove_spectral_norm(c2)
w3 = c3.weight.data.clone()
o3 = c3(t0)

c4 = torch.nn.utils.spectral_norm(c3)
w4 = c4.weight.data.clone()
o4 = c4(t0)

print(f"(o1==o2).float().mean(): {torch.eq(o1,o2).float().mean()}")
print(f"(o2==o3).float().mean(): {torch.eq(o2,o3).float().mean()}")
print(f"(o3==o4).float().mean(): {torch.eq(o3,o4).float().mean()}")
print(f"(o1==o3).float().mean(): {torch.eq(o1,o3).float().mean()}")
print(f"(o2==o4).float().mean(): {torch.eq(o2,o4).float().mean()}")
print()
print(f"(w1==w2).float().mean(): {torch.eq(w1,w2).float().mean()}")
print(f"(w2==w3).float().mean(): {torch.eq(w2,w3).float().mean()}")
print(f"(w3==w4).float().mean(): {torch.eq(w3,w4).float().mean()}")
print(f"(w1==w3).float().mean(): {torch.eq(w1,w3).float().mean()}")
print(f"(w2==w4).float().mean(): {torch.eq(w2,w4).float().mean()}")
</code></pre>
<p>Essentially:</p>
<ul>
<li>c1, c2, c3, and c4 are the same conv layer, but with SN added/removed several times</li>
<li>o1,o2,o3,o4 correspond to the output of these layers with the same input</li>
<li>w1,w2,w3,w4 correspond to the weights of the layers</li>
</ul>
<p>The output of the above script is typically as follows:</p>
<pre><code class="lang-auto">(o1==o2).float().mean(): 0.0
(o2==o3).float().mean(): 1.0
(o3==o4).float().mean(): 0.0533749982714653
(o1==o3).float().mean(): 0.0
(o2==o4).float().mean(): 0.0533749982714653

(w1==w2).float().mean(): 1.0
(w2==w3).float().mean(): 0.0
(w3==w4).float().mean(): 1.0
(w1==w3).float().mean(): 0.0
(w2==w4).float().mean(): 0.0
</code></pre>
<p>My understanding of SN is not very deep, and I do not fully understand the math behind it, but I think it is roughly like so:</p>
<ul>
<li>Spectral Norm enforces a constraint on the weights of the layer it is applied on, forbidding some possible values</li>
<li>This constraint is at least partially implemented in Pytorch by decomposing the weight matrix of the affected layer into weight, weight_u, and weight_v somehow</li>
</ul>
<p>Given that, here is what I think of the results:</p>
<ol>
<li>o1!=o2: this seems logical, since c1 was not subject to any restrictions, therefore it is to be expected that applying SN can change the layer’s behavior</li>
<li>o2==o3: this seems logical, since lifting the SN restriction should not modify the functionality of the layer</li>
<li>o3!=o4: <strong>this is unexpected</strong>, since c2 should have fit the SN restriction, and c3 represents a function that is equivalent to c2… c3 should therefore also fit the SN restriction without modification… or so I thought. Interestingly, running the same script several times actually sometimes give me o3==o4, which would fit my expectations… but not always.</li>
<li>o1!=o3: this follows from 1</li>
<li>o2!=o4: <strong>this is what prompted me to investigate</strong>, and it seems the issue comes from 3</li>
<li>w1 == w2: <strong>not exactly what I expected, but understandable</strong>. Looks like the SN restriction is fulfilled by doing extra steps during the convolution using weight_u and weight_v, but weight is kept as it was.</li>
<li>w2!=w3: logical, since we need o2==o3, weight_u and weight_v must be directly incorporated back into weight at this step</li>
<li>w3==w4: logical, this is just a repeat of 6</li>
<li>w1!=w3: this follows from 7</li>
<li>w2!=w4: not exactly what I expected, but understandable given 7</li>
</ol>
<p>Is my understanding correct?<br>
<strong>What method can I use to temporarily remove SN from my model, and then put it back, without it messing up the already learned weights?</strong></p>