[feature request] Spectral norm support in torch.norm or factor out Power Iteration from spectralnormalization in some other place and orthogonalization from PowerSGD hook

[vadimkantorov]

Currently PyTorch has a spectral matrix norm implementation based on power iteration hidden inside the SpectralNorm. Since new torch.norm supports matrix norms (e.g. nuclear norm), it may be good to move it there, so it's usable in other contexts too (e.g. for debugging/tracing spectral properties of neural net weights, and since PyTorch doesn't support truncated randomized SVD yet).

[crcrpar]

I like the idea.

What API do you assume after supports spectral_norm using power iteration method which requires 2 vectors.

[vadimkantorov]

NumPy does this by supporting additional p= arguments:

2 | 2-norm (largest sing. value)
-2 | smallest singular value

[crcrpar]

Ah, I'm concerned with how to handle u and v.
Will the norm function take u as input with dim or prepare inside it?

[vadimkantorov]

• may be good to factor out orthogonalization out of powersgd and put it in a proper function: [Proposal] Use batched oprations to accelerate PowerSGD #74907

[tvogels]

@vadimkantorov, factoring out orthogonalization from PowerSGD is a good idea, but the orthogonalization operation there is currently much slower than it could be. Based on recommendations in the appendix of this paper (and with help from one of the authors), @younik is currently looking into replacing it with a custom CUDA kernel for the Householder method. It is work in progress, but we are seeing significant speedups. We hope to push a PR for this very soon.

Speedups

Matrix shape   Custom kernel (μs)   torch.qr (μs)   Speedup
(2, 512)       67.4                 11176           165.9x
(2, 4096)      96.5                 11774           122.0x
(4, 512)       96.3                 10497           109.0x
(4, 4096)      71.2                 12448           174.8x
(8, 512)       70.8                 12177           172.1x
(8, 4096)      88.1                 12500           141.9x
(16, 512)      80.2                 11780           147.0x
(16, 4096)     164.9                13296           80.6x
(32, 512)      148.1                11872           80.1x
(32, 4096)     319.1                18132           56.8x
(64, 512)      291.0                11562           39.7x
(64, 4096)     678.6                12940           19.1x
(128, 512)     573.4                16406           28.6x
(128, 4096)    1299.5               27551           21.2x

[vadimkantorov]

PR on orthogonalize: #76673

[redwrasse]

Another consideration in broadening the applicability of spectral norm is either in documentation or implementation, control of convergence properties- for spectral norm and sister operations. For spectral norm, I believe it's tied to the ratio of top two eigenvalues for the power iteration method. I found today some torch documentation discussing how to handle such concerns in the linalg package: https://pytorch.org/docs/stable/notes/numerical_accuracy.html#extremal-values-in-linalg

Spectral operations like svd, eig, and eigh may also return incorrect results (and their gradients may be infinite) when their inputs have singular values that are close to each other. This is because the algorithms used to compute these decompositions struggle to converge for these inputs.