KATE: https://arxiv.org/pdf/1705.02033.pdf

k_competitive_layer.py

def k_competitive_layer(tensor, top_k, alpha, cuda=False):
    """
    https://arxiv.org/pdf/1705.02033.pdf
    
    Args:
        tensor: Tensor
        top_K: int
        alpha: hyper param
    """
    def k_update(tensor_, indices, top_k, alpha, largest=True, cuda=False):
        """
        """
        tensor = tensor_.clone()
        
        # Do we have enough top ks on each row
        # to perform topk
        has_enough_top_ks = torch.sum(indices, 1)

        for i in range(has_enough_top_ks.size(0)):
            if has_enough_top_ks[i].data[0] > top_k:                
                # Current tensor and indices
                c_tensor = tensor[i]
                
                if largest:
                    c_indices = c_tensor > 0
                else:
                    c_indices = c_tensor < 0

                # get top k/2 neg/pos activations                             
                topk_val, topk_indices = torch.topk(c_tensor, int(top_k/2), dim=0, largest=largest)                

                # Get energy of lower (k/2)
                E_ = torch.sum(c_tensor[c_indices]) - torch.sum(topk_val)
                E_ = E_.data[0]

                # Create a matrix to store energies
                energies = Variable(torch.zeros(c_tensor.size()))        
                energy_var_placeholder = Variable(torch.zeros(topk_indices.size()))

                if cuda:
                    energies = energies.cuda()
                    energy_var_placeholder = energy_var_placeholder.cuda()                    

                energies = energies.scatter(0, topk_indices, energy_var_placeholder + (E_*alpha))

                # Create mask to zero non top k activations
                zero_mask = c_indices.scatter(0, topk_indices, energy_var_placeholder.byte())

                # Zero out non top-k neurons
                c_tensor = c_tensor.masked_fill(zero_mask, 0)

                # Add energies            
                tensor[i] = c_tensor + energies
        
        return tensor
    
    # Get indices
    pos_indices = tensor > 0        
    tensor = k_update(tensor, pos_indices, top_k, alpha, largest=True, cuda=cuda)
    
    neg_indices = tensor < 0
    tensor = k_update(tensor, neg_indices, top_k, alpha, largest=False, cuda=cuda)
        
    return tensor


class KateAutoencoder(nn.Module):
    def __init__(self, n_tokens, n_hidden, top_k, alpha=1):
        super().__init__()
        assert (n_tokens > top_k)
        
        self.alpha = alpha        
        self.top_k = top_k
        
        self.n_tokens = n_tokens
        
        self.fc1 = nn.Linear(n_tokens, n_hidden)
        self.fc2 = nn.Linear(n_hidden, n_tokens)
    
    def forward(self, x):        
        z = F.tanh(self.fc1(x))
        z_hat = k_competitive_layer(z, self.top_k, self.alpha, cuda=CUDA)
        x_hat = F.sigmoid(self.fc2(z_hat))
        return x_hat