Neuralized K-Means

[jackmcrider]

Hi chr5tphr!

I started an attempt to implement (deep) neuralized k-means (https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=9817459) as more people want to use it and ask for code.

I took the SoftplusCanonizer from the docs as a starting point.

Main changes:

• Add Distance, NeuralizedKMeans and LogMeanExpPool in zennit.layer
• Add KMeansCanonizer in zennit.canonizers

Some things can be optimized:

• add tests (I wrote a few, but then changed something and did not bother to rewrite tests unless I see a chance of this being merged)
• copy.deepcopy in KMeansCanonizer.register probably not ideal
  • alternative: clone distance_module.centroids and import Distance from layer.py; then, create new Distance instance in remove()
• KMeansCanonizer.apply scans for Distance layers, but there can be distance layers in some architectures that don't belong to k-means
• one idea would be add a "contrastive layer" and identify kmeans as a composition of Distance followed by Contrastive similar to MergeBatchNorm:
  • contrastive computes (out[:,None,:] - out[None,:,:])[mask].reshape(K,K-1,D), cf. line 379-384 in canonizers.py
• several advantages:
  • output of k-means and neuralized k-means are identical,
  • contrastive layer could also be applied to classifiers to get class-contrastive explanations (cf. Fig. 10.5 in https://iphome.hhi.de/samek/pdf/MonXAI19.pdf)
  • could replace NeuralizedKMeans layer by something more general (difference between squared distances and difference between linear layers are both linear, i.e. a contrastive layer covers both; would still need two separate Canonizers)
• the scanning in KMeansCanonizer.apply does not work if Distance is the first layer (kmeans in input space), one can do Sequential(Distance(centroids)) as a trick, but not ideal
• if I understand correctly, KMeansCanonizer.register is executed each time one calls with Gradient(...) as attributor; could be a bottleneck if number of clusters is large

Closes #198

[jackmcrider]

I just found the contributing guide and converted this to draft for now since I broke every single guideline.

@p16i I clicked somewhere and triggered a review request. Please ignore.

[jackmcrider]

I tried to merge everything into one commit, extended documentation and made sure that all checks pass.
Could be reviewed, but there are no functional changes.

[jackmcrider]

Checks pass 👍

It's quite challenging to get reproducible tox results for tutorials (e.g. had to manually fiddle with metadata.kernelspec.name in the raw .ipynb before committing). Probably a limitation of tox, but could be documented for future contributors. I'm not sure where, maybe in Contributing#continuous-integration.

I'm gonna freeze this branch for now, unless something comes up.

[chr5tphr]

Hey @jackmcrider

thanks a lot for the contribution!
I have looked at your implementation and left a few comments.

I have not yet looked at the tutorial.

[jackmcrider]

I have commited a new version with roughly these changes:

• remove LogMeanExpPool
• add layers MinPool1d and MinPool2d (simple inheritance from the PyTorch MaxPool* classes)
• add rules MinTakesMost1d, MinTakesMost2d, MaxTakesMost1d, MaxTakesMost2d
• change KMeansCanonizer to use MinPool1d instead of LogMeanExpPool
• change tutorial to use MinTakesMost1d at the output layer
• apply changes requested in review (rename Distance to PairwiseCentroidDistance, do self.parent_module.add_module instead of setattr(self.parent_module...)

I'm not sure if we want four rules for the *TakesMost* or if one rule with mode='max'/mode='min' and some autodetection for the 1d/2d is better.