Add exact spectral norm feature #831
Conversation
|
Open to discussion on best interface. The current caution around bias is ok but could be a footgun (although the original paper doing the reshaping is already kind of a footgun imo lol), but maybe I can apply the linear transpose around the whole layer pytree and update the bias too? I'm not sure that would 100% generalize to all layers though. |
| !!! Caution | ||
|
|
||
| If `exact` is true, it computes the transpose via `jax.linear_transpose` of | ||
| the layer. This includes all operations of the layer call, which means for | ||
| layers with a bias, this can result in the incorrect spectral value. |
There was a problem hiding this comment.
Maybe we should follow JAX's lead here and transpose the tangent pass of jax.jvp?
There was a problem hiding this comment.
Hmmm, I tried to implemented what I thought you meant. This also means we could remove the "weight" flag for the exact case (maybe?) since we basically "determine" the weight through the jvp?
Let me know if this is what you had in mind, or if I was totally off. Does seem like a lot of jvps tho.
Addresses part of #810.