Differentiable Permutation

username: nreHieW
cooked: <p>Hello, I was wondering if there is a way to perform a differentiable reordering of items such that the model learns the optimal ordering of items. Specifically, I have an input tensor of shape <code>[batch_size, seq_length, dim]</code> from a VIT block and I want to shuffle the items along the second dimension such that the output tensor is of the same shape just with the items in a different order.</p>
<p>This is what I have so far</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
from fast_soft_sort.pytorch_ops import soft_rank # differentiable ranking taken from https://github.com/google-research/fast-soft-sort/tree/master

class Shuffle(nn.Module):
    def __init__(self, L):
        """
        L: Length of the sequence
        D: Dimension of each vector in the sequence
        """
        super(Shuffle, self).__init__()
        self.weights = nn.Parameter(torch.randn(L)).unsqueeze(0) * 100

    def forward(self, x):
        batch_size, seq_length, dim = x.shape
        expanded_scores = self.weights.unsqueeze(-1)
        indices = soft_rank(self.weights).squeeze(0).long() - 1 # detach here
        indices = indices.view(1, -1, 1).expand_as(x)
        x = torch.gather(x, dim = 1, index = indices)
        return x
</code></pre>
<p>The idea I had was for the model to learn <code>self.weights</code>, use a differentiable ranking algorithm and then gather these ranked items as indices. The problem is that casting to <code>long()</code> detaches gradients.</p>
<p>I came across this <a href="https://discuss.pytorch.org/t/using-torch-gather-properly/158372">thread</a> for selecting rather than shuffling. I am unsure if its possible to apply that to my use case.</p>
<p>All help is greatly appreciated, thanks!</p>

username: J_Johnson
cooked: <p>In order to implement something like this, here is what I would do:</p>
<ol>
<li>Have some non-model way of determining what is optimal. This could be accomplished by testing all possible permutations and having some pre-determined algorithm for scoring them based on whatever parameters you desire, or by manually scoring them via whatever you consider to be optimal. With this, you can designate targets. Targets should be from 0 to n-1 where n is the possible types of permutations.</li>
<li>Create a convolution network of appropriate dimensions for the inputs. Output of the model should be of size n classes.</li>
<li>Train the model on the inputs and targets.</li>
</ol>
<p>This way, you can just treat it as a classification problem using cross entropy loss with logits between the outputs and targets.</p>
<p>Aside from gradients being detached when changing to <code>.long()</code> issue, the model that you have proposed has no means of determining the loss. Also, in order to get some type of classification loss, you want an output of logits or probabilities.</p>

username: nreHieW
cooked: <p>Thanks for the reply. This was meant to be part of a ViT so<br>
ViT: [shuffle → Attn ] x N times.</p>
<p>Thats how I wanted to implement the model. Is there a solution that can fit this use case?</p>

username: J_Johnson
cooked: <p>I would encourage you to have a look at some reinforcement learning models.</p>
<p>At any rate, you want the “Shuffle model” to take the same input as the ViT, and reduce it down to classes. Convolution operations are your best bet.</p>
<p>The class output would be the permutation applied. After the permutation is made on the original image, pass that into the ViT. Then get the loss of the ViT. Set the target classes to 0s(assuming logits as outputs) and then alter the class choice actually made to be (1 - ViT_loss). Keep in mind, you need to get the loss before the mean is taken across the batch, and detach a copy of this from the ViT. The batches with lower loss, will get a stronger signal for that choice fed back as a target.</p>
<p>Use l1loss between the ‘Shuffle model’ outputs and the targets.</p>
<p>Additionally, you’ll want a greedy epsilon function that randomly chooses for the ‘Shuffle model’ initially. This ensures it explores all permutations. The greedy epsilon function should decay over time. See a DQN tutorial for examples of this.</p>
<p>You could also experiment with providing some normalization to the reward. Instead of (1 - ViT_loss), could do (ViT_loss_max - ViT_loss) or just normalize per batch: 1-(ViT_loss - ViT_loss_min)/(ViT_loss_max - ViT_loss_min) with an if statement that skips when ViT_loss_max == ViT_loss_min.</p>