Playing with alpha, I found very strong results with 1 shared alpha value across the layer. Initial observations suggest that one alpha per kernel doesn't add anything to performance, while incurring the addition of extra parameters.

Inspecting the alpha values chosen by LSUV_like_init(), they fall roughly on the range from 10-50 for a small model I'm running on Fashion MNIST. Choosing an initial alpha of 10 and disabling auto init generates the strongest results I've seen from my model yet.

@ducha-aiki, what do you think about moving to a single alpha per layer and removing the LSUV init? You've spent the most time thinking about it, so it's possible you've built an intuition around it I don't have yet.

@4rtemi5 based on what I've seen, the alpha boosts performance noticeably. Based on what implementation we settle on in PyTorch I plan to mimic it in the Jax code. It might make a difference in getting to 90% on CIFAR-10.

This diff is a little messy-it also has some changes to the demo scripts to make them compatible with the new API.

This PR proposes a change in implementation of SCS which uses Conv2D, a much more performant primitive. The gist of the idea is to

1. Normalize the kernel as we already do
2. Compute per-window normalization factor coming from the image by sqrt(avg_pool2d(input ** 2) * window_size)
3. Compute the dot products via conv2d between the input and normalized kernel
4. Normalize the output elementwise by the array obtained in step 2.
5. Proceed as before

On my laptop this achieves more than 2x performance improvement. I attach a quick test (compat_test.py) to verify that the two implementations yield (almost) the same results. This PR will obviously need some code unification and renaming before being merged, I submit it as it to facilitate easy comparisons by the maintainers :)

Hello, I am unable to understand the testing loss. Why loss = F.cross_entropy(preds, labels) is not also used for test loss with test preds?

https://github.com/brohrer/sharpened_cosine_similarity_torch/blob/28c135adfb6062b0e9e29bb906736ceea9c08162/model_cifar10_18_4.py#L152-L163

Debug and test a maximum magnitude pooling layer

@StephenHogg wrote a nice looking bit of code for this, but I haven't been able to get it to work just yet.

Original code https://github.com/StephenHogg/SCS/blob/main/SCS/layer.py

Copy/pasted to https://github.com/brohrer/sharpened_cosine_similarity_torch/blob/main/absolute_pooling.py

Updated PyTorch and Jax implementations Removed p parameter exponentiation Added weight clipping and p parameter lower limit Changed parameter initialization Removed alpha parameter

Demo scripts and MNIST test all run.

Any thought or comments? @ducha-aiki @4rtemi5

This is an example application of the scsl fashion mnist example with PyTorch lightning. Lightning removes a lot of the boilerplate of PyTorch and helps to focus on the network. Furthermore, lightning makes training on GPUs very simple.

I am a physician and don't know what I am doing, but this helps to prevent the NaNs, maybe it helps you to create a more sophisticated solution.

        out = inp.square()
        if torch.any(torch.isnan(out)):
            raise Exception("out")
            
        if self.groups == 1:
            out = out.sum(1, keepdim=True)

        norm = F.conv2d(
            out,
            torch.ones_like(self.weight[:1, :1]),
            None,
            self.stride,
            self.padding,
            self.dilation) + 1e-6
        
        if torch.any(torch.isnan(norm)):
            raise Exception("norm")
            
        # prevent 0 and inf
        q = torch.exp(-self.q / (self.q_scale**2 + 0.1))
        if torch.any(torch.isnan(q)):
            raise Exception("q")
            

        weight = self.weight / (self.weight.square().sum(dim=(1, 2, 3), keepdim=True).sqrt() + q + 1e-6)
        if torch.any(torch.isnan(weight)):
            raise Exception("weight")
            

        out = F.conv2d(
            inp,
            weight,
            self.bias,
            self.stride,
            self.padding,
            self.dilation,
            self.groups
        ) / ((norm**2).sqrt() + 1e-1)

        if torch.any(torch.isnan(out)):
            raise Exception("out2")

        # Comment these lines out for vanilla cosine similarity.
        # It's ~200x faster.
        abs = (out.square() + 1e-6).sqrt()
        sign = out / abs
        # prevent 0 and inf
        p = torch.exp(self.p / (self.p_scale**2 +0.1))
        out = abs ** p
        out = out * sign
        if torch.any(torch.isnan(out)):
            raise Exception("out3")
        return out```

Hi! I played around with your code yesterday and ended up refactoring it quite heavy to ease usage. The code is nowhere ready for any pull request, but I hope that you find the changes somewhat useful.

You can find the code here.

Notes:

• I were able to achieve 91.3% CIFAR-10 accuracy with a model of 1.2M parameters( see here for model implementation).
• Residual connections seem to somewhat help deeper networks, however, training a much deeper model does not work out of the box (see here for implementation).
• Either annotating the forward pass of the SCS layer or scripting the model will improve runtime at least 30% (measured on NVIDIA 1060 GPU).

The code automatically uses a CUDA capable device if available, is quite cleaner than the original training script and will hopefully help you iterate on the idea :)

You can view the tensorboard logs of my experiments here: https://tensorboard.dev/experiment/AY27LbxrRpaBHNMO0m9Wkw/#scalars

Issue due to progress bar after every epoch it prints something like,

0it [00:00, ?it/s]run: 0   epoch: 0   duration: 75.29   training loss: 2.249   testing loss: 2.137   training accuracy: 40.17%   testing accuracy: 57.66%

Instead of,

run: 0   epoch: 0   duration: 75.29   training loss: 2.249   testing loss: 2.137   training accuracy: 40.17%   testing accuracy: 57.66%

So far sharpened cosine similarity has performed well on Fashion MNIST. The next step is to demonstrate it on CIFAR 10.

Here are some partial implementations in case they help from Abhiraj Kanse https://colab.research.google.com/drive/1lKx9wiY2bMg5rH5EvJKymj7nb8I9ESc4?usp=sharing

from @oliver-batchelor https://github.com/oliver-batchelor/scs_cifar/blob/main/src/main.py

Enhancement

• Needs a progress bar at least in train loader.

Issues

Multiple files have the following issues.

• Duplicate calculation of epoch time inside mini batch loader. Extra calculation of epoch_duration = time.time() - epoch_start_time inside train loader.

https://github.com/brohrer/sharpened_cosine_similarity_torch/blob/690909ff360c67581ab2745a8a8516c4ee133ad8/demo_fashion_mnist.py#L127-L131

• Test preds calculation not used and can be removed.

https://github.com/brohrer/sharpened_cosine_similarity_torch/blob/690909ff360c67581ab2745a8a8516c4ee133ad8/demo_fashion_mnist.py#L138-L147

• Make following consistent,

https://github.com/brohrer/sharpened_cosine_similarity_torch/blob/690909ff360c67581ab2745a8a8516c4ee133ad8/model_cifar10_18_4.py#L127-L130

https://github.com/brohrer/sharpened_cosine_similarity_torch/blob/690909ff360c67581ab2745a8a8516c4ee133ad8/model_cifar10_18_4.py#L154-L155