
Of course, this is an oversimplification, and there won't exist a unique stencil that perfectly represents all of the contributions of the input to a class (otherwise that would mean that the problem could be solved linearly). So this should be used for visualization purposes only, not as an accurate predictor.
Once we have a set of stencils for each class, we can effortlessly measure input contribution without relying on any external visualization techniques.
To find a stencil, we can start from a "random noise" stencil and optimize it to maximize the activations for a specific class by just back-propagating through the MLP-Mixer.



By doing this we can end up with many valid stencils, and we can reduce them to a few by using K-means to cluster them into similar stencils and averaging each cluster.
Using the Mixer to get the best of each world