Multi-head attention and grouped convolutions

[PythonNut]

As we discussed previously, the abstractions of our current weight matching algorithm do not generalize to networks with the following structure:

1. There are n separate things which can be permuted with each other.
2. Each of those n things contains k units which can be permuted.

Probably the most pressing example of this is multi-head attention (where the heads can be permuted and also each head has its own latent permutation). Grouped convolutions are another example.

If our cost is linear, we can apply a hierarchical matching algorithm which can match these local structures exactly, however it's expensive (runs in O(n²k³ + n³) time).

For applications like fine-tuning where the weights are already quite similar, it may be possible to ignore the outermost permutation and still achieve good results.