Spectral Norm Features

[lockwo]

The original spectral norm paper didn't accurately estimate spectral norms (e.g. in the convolution case they just reshape it then do normal matrix, but this isn't the actual transpose operation, see: https://arxiv.org/abs/2009.02773 for more) which is now the general default in neural network packages (equinox included). However, it's not that much LoC to get a spectral norm that can do one or both of the following (sources: https://arxiv.org/abs/1804.04368, https://arxiv.org/abs/1811.07457)

• actually approximates the correct spectral norm (and thus adheres to the 1-lipshitz cap, which original SN isn't super great at, see fig 3 of https://arxiv.org/abs/1802.04034). This basically just replaces the reshaping -> transpose, with the actual linear transpose operation of a given layer
• can be made numerically meaningful (by guaranteeing convergence to a certain threshold), at the cost of only being forward differentiable

I have some code that does both of these things that I could polish enough for equinox to really boost the spectral norm flexibility (of course keeping all existing features, these are valuable in research, but in actual performance doing a poor approximation to a spectral norm actually often performs better than the real norm lol see sec 4 of http://www.arxiv.org/abs/2009.02773). But I wanted to check if there was any interest/scope for these sort of enhancements/expansions.

[patrick-kidger]

Definitely in-scope! :)