How is spectral norm of a parameter calculated?
when I do,

import torch, torch.nn as nn
x = nn.Linear(3, 3)
y = torch.nn.utils.spectral_norm(x)


then it gives four different weight matrices,

y.weight_u

tensor([ 0.6534, -0.1644,  0.7390])


y.weight_orig

Parameter containing:
tensor([[ 0.2538,  0.3196,  0.3380],
        [ 0.4946,  0.0519,  0.1022],
        [-0.5549, -0.0401,  0.1654]], requires_grad=True)


y.weight_v

tensor([-0.3650,  0.2870,  0.8857])


y.weight

tensor([[ 0.5556,  0.6997,  0.7399],
        [ 1.0827,  0.1137,  0.2237],
        [-1.2149, -0.0878,  0.3622]], grad_fn=&lt;DivBackward0&gt;)


how are these four matrices calculated?

I just finished reading the paper for this method which can be found on arxiv. If you have the appropriate mathematical background I would recommend reading it. See appendix A for the power algorithm which describes what u and v are.

That said I'll try to summarize here.

First, you should know that the spectral norm of a matrix is the maximum singular value. The authors propose finding the spectral norm of weight matrix W, then dividing W by its spectral norm to make it close to 1 (justification for this decision is in the paper).

While we could just use torch.svd to find a precise estimate of the singular values, they instead use a fast (but imprecise) method called "power iteration". Long story short, the weight_u and weight_v are rough approximations of the left and right singular vectors corresponding to the largest singular value of W. They are useful because the associated singular value, i.e. the spectral norm, of W is equal to u.transpose(1,0) @ W @ v if u and v are the actual left/right singular vectors of W.


y.weight_orig contains the original values in the layer.
y.weight_u is the approximation of the first left singular vector of y.weight_orig.
y.weight_v is the approximation of the first right singular vector of y.weight_orig.
y.weight is the updated weight matrix which is y.weight_orig divided by its approximate spectral norm. 


We can verify these claims by showing that the actual left and right singular vectors are nearly parallel to y.weight_u and y.weight_v

import torch
import torch.nn as nn

# pytorch default is 1
n_power_iterations = 1

y = nn.Linear(3,3)
y = nn.utils.spectral_norm(y, n_power_iterations=n_power_iterations)

# spectral normalization is performed during forward pre hook for technical reasons, we
# need to send something through the layer to ensure normalization is applied
# NOTE: After this is performed, x.weight is changed in place!
_ = y(torch.randn(1,3))

# test svd vs. spectral_norm u/v estimates
u,s,v = torch.svd(y.weight_orig)
cos_err_u = 1.0 - torch.abs(torch.dot(y.weight_u, u[:, 0])).item()
cos_err_v = 1.0 - torch.abs(torch.dot(y.weight_v, v[:, 0])).item()
print('u-estimate cosine error:', cos_err_u)
print('v-estimate cosine error:', cos_err_v)

# singular values
actual_orig_sn = s[0].item()
approx_orig_sn = (y.weight_u @ y.weight_orig @ y.weight_v).item()
print('Actual original spectral norm:', actual_orig_sn)
print('Approximate original spectral norm:', approx_orig_sn)

# updated weights singular values
u,s_new,v = torch.svd(y.weight.data, compute_uv=False)
actual_sn = s_new[0].item()
print('Actual updated spectral norm:', actual_sn)
print('Desired updated spectral norm: 1.0')


which results in

u-estimate cosine error: 0.00764310359954834
v-estimate cosine error: 0.034041762351989746
Actual original spectral norm: 0.8086231350898743
Approximate original spectral norm: 0.7871124148368835
Actual updated spectral norm: 1.0273288488388062
Desired updated spectral norm: 1.0


Increasing the n_power_iterations parameter will increase the accuracy of the estimate at the cost of computation time.
