Pytorch implementation of the outlier suppression loss described in the paper "Improving generalization by loss modification" by Michael Tetelman

outlier_supression_loss.py

# Pytorch implementation of the outlier suppression loss described in the paper
# "Improving generalization by loss modification" by Michael Tetelman 
# https://openreview.net/forum?id=vHOO1lxggJ

from torch import nn
import torch.functional as F

def outlier_supression_loss(input, target):
    return F.softplus(F.nll_loss(F.log_softmax(input, dim=-1), target, reduction='none')).mean()

class OutlierSuppressionLoss(nn.Module):
    def forward(self, input, target):
        return outlier_supression_loss(input, target)