I was planning on doing the first idea, and I do like the second idea! I'm slightly skeptical that the two rotations will be the same, but I did find that when performing the method on the last layer of the model, I get the identity matrix, which is some evidence in favor of the 'rotations are the same' prediction being right in general.

Suppose our  model has the following format:

$\text{Model}\left(\text{input}\right)=(M_3\circ NL\circ M_2\circ NL\circ M_1)\left(\text{input}\right)$

where $M_3,M_2,M_1$ are matrix multiplies, and $NL$ is our nonlinear layer.

We also define a sparsity measure to minimize, chosen for the fun property that it really really really likes zeros compared to almost all other numbers.

$\text{Sparsity}\left(A\right)=-{\sum }_{i,j}\frac{1}{|a_{i,j}|+0.1}$

note that lower sparsity according to this measure means more zeros.

There are two reasonable ways of finding the right rotations. I will describe one way in depth, and the other way not-so in depth. Do note that the specifics of all this may change once I run a few experiments to determine whether there's any short-cuts I'm able to take^[1].

We know the input is in a preferred basis. In our MNIST case, it is just the pixels on the screen. These likely interact locally because the relevant first-level features are local. If you want to find a line in the bottom right, you don't care about the existence of white pixels in the top left.

We choose our first rotation $R_1$ so as to minimize 

$\text{Sparsity}\left(R_1{J}_1\right)$

where

$J_1:=J_{\text{input}}(NL\circ M_1)$

Then the second rotation $R_2$ so as to minimize

$\text{Sparsity}\left(R_2{J}_2\right)$

where

$J_2:=J_{(R_1\circ NL\circ M_1)\left(\text{input}\right)}(NL\circ M_2\circ R_1^{-1})$

And finally choosing $R_3$ so as to minimize 

$\text{Sparsity}\left(R_3{J}_3\right)$

where

$J_3:=J_{(R_2\circ NL\circ M_2\circ NL\circ M1)\left(\text{input}\right)}(M_3\circ R_2^{-1})$.

The other way of doing this is to suppose the output is in a preferred basis, instead of the input.

Currently I'm doing this minimization using gradient descent (lr = 0.0001), and parameterizing my rotation matrices using the fact that if $A$ is an antisymmetric matrix^[2], then $e^A$ is a rotation matrix, and that you can make an antisymmetric matrix by choosing any old matrix $B$, then doing $A:=B-B^T$. So we just figure out which $B$ gets us an $R=e^{B-B^T}$ which has the properties we like.

There is probably a far, far better way of solving this, other than gradient descent. If you are interested in the specifics, you may know a better way. Please, please tell me a better way!

1. ^{^}

   An example of a short cut: I really don't want to find a rotation which minimizes average sparsity across every input directly. This sounds very computationally expensive! Does minimizing my sparsity metric on a particular input, or only a few inputs generalize to minimizing the sparsity metric on many inputs?

2. ^{^}

   Meaning its a symmetric matrix with its top right half the opposite sign as it's bottom left half.

What do you think about the Superposition Hypothesis? If that were true, then at a sufficient sparsity of features in the input there is no basis in which the network is thinking in, meaning it will be impossible to find a rotation matrix that allows for a bijective mapping between neurons and features.

I'd say that there is a basis the network is thinking in in this hypothetical, it would just so happens to not match the human abstraction set for thinking about the problem in question.

If due to superposition, it proves advantageous to the AI to have a single feature that kind of does dog-head-detection and kind of does car-front-detection, because dog heads and car fronts don't show up in the training data at the same time, so it can still get perfect loss through a properly constructed dual-purpose feature like this, it'd mean that to the AI, dog heads and car fronts are "the same thing". 

The network hasn't figured out how to distinguish between them. In a more general data set where dog heads and car fronts can co-occur, this network would fail. Its abstractions are optimised for the narrow training data set, where it genuinely proved to be unnecessarily cumbersome to assign different concepts to those two things.

As AIs get more capable and general, I'd expect the concepts/features they use to start more closely matching the ones humans use in many domains. As AI gets superhuman, I would be somewhat worried about it finding new concept/feature sets that work even better and more generally than human ones.

I mostly expect networks at zero loss to not to be in a superposition, since we should expect those networks to be in a broad basin, meaning fairly few independent, orthogonal, features, so less room to implement two completely different functions. But we don't always find networks in broad basins, so we may see some networks in a superposition.

It would be interesting to study which training regimes and architectures most/least often produce easily-interpretable networks by this metric, and this may give some insight into when you see superposition.

In the cases where there is a nice basis this device finds, we may also expect it to disentangle any superpositions which exist, and for this superposition to be a combination of two fairly simple functions, requiring very few features, or interpreting the same features in different ways.

Think of the rotation as an adjustment to our original program's representation, which keeps the program the same, but hopefully makes it clearer what it's doing. It would probably be interesting to see what exactly the adjustment is doing, like how it's interesting to see how someone would go about translating assembly to a python script, but not particularly your highest priority when trying to figure out what the assembly program is doing after you have the corresponding python script.

You can easily do the same with any affine transformation, no? Skew, translation (scale doesn't matter for interpretability).

You can do this with any normalized, nonzero, invertible affine transformation. Otherwise, you either get the 0 function, get a function arbitrarily close to zero, or are unable to invert the function. I may end up doing this.

What would suffice as convincing proof that this is valuable for a task: the transformation increases the effectiveness of the best training methods.

This will not provide any improvement in training, for various reasons, but mainly because I anticipate there's a reason the network is not in the interpretable basis. Interpretable networks do not actually increase training effectiveness. The real test of this method will be in my attempts to use it to understand what my MNIST network is doing.