I am trying to implement the “soft-softmax” loss for multi-label classification that is very briefly described on page 5 (loss function) of https://arxiv.org/abs/1805.00932.
I have an implementation where I computed the gradients myself and implemented a custom loss function with forward-backward (see CategoricalMultiLabelLoss below). I used this blogpost 3 as a reference.
To simplify, I tried to compare this with computing the loss using pytorch operators and let autograd do the work. While the forward passes yield identical loss value, the parameter gradients when I do a backward pass are different.
Can someone help me understand why this might be the case and what I am missing? Is the loss non-differentiable at certain points due to which autograd might not work well? (I didn’t think so but may be I’m wrong)
import torch
import torch.nn as nn

# custom loss implementation
class CategoricalMultiLabelLoss(torch.autograd.Function):  
    @staticmethod
    def forward(ctx, inputs, targets):
        scores = inputs.data
        batch_size = len(scores)
        # Compute cross-entropy loss
        scale_factor = (targets > 0).sum(dim=1, keepdim=True).float()
        logprobs = (scores * targets).div(scale_factor)
        data_loss = - torch.sum(logprobs)/ batch_size
        ctx.save_for_backward(inputs, targets)
        return data_loss

    @staticmethod
    def backward(ctx, grad_output):
        inputs, targets = ctx.saved_tensors
        scale_factor = (targets > 0).sum(dim=1, keepdim=True).float()
        scale_factor_repmat= scale_factor.repeat(1,targets.size(1))
        print(scale_factor.shape, scale_factor_repmat.shape, scale_factor_repmat)
        delta = inputs.data.exp()   # If the class label is 0, the gradient is equal to probs
        delta = scale_factor_repmat *(delta-1) + (1-scale_factor_repmat)*delta
        delta = delta * targets
        return delta, None
    
# dummy model
class DummyModel(nn.Module):
    def __init__(self, num_inputs, emb_size, emb_dropout, hidden_size, num_outputs, pad_index):
        self.num_inputs = num_inputs
        self.emb_size = emb_size
        self.emb_dropout = emb_dropout

        self.hidden_size = hidden_size
        self.num_outputs = num_outputs
        self.pad_index = pad_index

        super(DummyModel, self).__init__()

        self.emb = nn.Embedding(num_inputs, emb_size, padding_idx=pad_index)
        self.emb_dropout = nn.Dropout(p=emb_dropout)        
        self.fc1 = nn.Linear(emb_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_outputs)
        self.logsoftmax = nn.LogSoftmax()

    def forward(self, input_seq, length, masks, debug=False):
        embs = self.emb(input_seq)
        embs_dropout = self.emb_dropout(embs)
        avg = torch.div(
            torch.sum(embs_dropout, axis=1),
            length
        )
        out1 = self.relu(self.fc1(avg))
        out2 = self.fc2(out1)
        
        # constraint: do not predict masks so set these scores to large negative value
        out2.scatter_(1, masks, -1e5)
        return self.logsoftmax(out2)
    
# initialize two identical models with dropout set to zero
num_inputs = 10
emb_size = 2
emb_dropout = 0.0
hidden_size = 1
num_outputs = 5
PAD_INDEX = 0
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

ma = DummyModel(
    num_inputs,
    emb_size,
    emb_dropout,
    hidden_size,
    num_outputs,
    PAD_INDEX
)
ma.to(DEVICE)

mn = DummyModel(
    num_inputs,
    emb_size,
    emb_dropout,
    hidden_size,
    num_outputs,
    PAD_INDEX
)
mn.to(DEVICE)
mn.load_state_dict(ma.state_dict())

# dummy data
ex = (
    # inputs
    torch.cuda.LongTensor([
        [2, 3, 0],
        [4, 2, 1]
    ]),
    # targets
    torch.cuda.FloatTensor([
        [0, 0, 0, 0.5, 0.5],
        [0, 1.0, 0, 0, 0]
    ]),
    # lengths
    torch.cuda.LongTensor([
        [2],
        [3]
    ]),
    # masks
    torch.cuda.LongTensor([
        [2],
        [4]
    ])
)

# Use custom loss with backward impl
i1, t1, il1, m1 = ex
o1 = ma(i1, il1, m1)
loss1 = CategoricalMultiLabelLoss.apply(o1, t1)
print(loss1)
o1.retain_grad()
loss1.backward()
print(o1.grad)
for p in ma.parameters():
    print(p.grad)

i2, t2, il2, m2 = ex
o2 = mn(i2, il2, m2)
loss2 = torch.mean( # average CE across batch
    torch.div(
        torch.sum( # compute CE
            (-t2 * o2), dim=1
        ),
        (t2 > 0).sum(dim=1) # scaling factor
    )
)
print(loss2)
o2.retain_grad()
loss2.backward()
print(o2.grad)
for p in mn.parameters():
    print(p.grad)

This is the output:
tensor(1.2839, device='cuda:0', grad_fn=<CategoricalMultiLabelLossBackward>)
torch.Size([2, 1]) torch.Size([2, 5]) tensor([[2., 2., 2., 2., 2.],
        [1., 1., 1., 1., 1.]], device='cuda:0')
tensor([[-0.0000, -0.0000, -0.0000, -0.9301, -0.8291],
        [-0.0000, -0.8359, -0.0000, -0.0000, -0.0000]], device='cuda:0')
tensor([[0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.]], device='cuda:0')
tensor([[0., 0.]], device='cuda:0')
tensor([0.], device='cuda:0')
tensor([[0.],
        [0.],
        [0.],
        [0.],
        [0.]], device='cuda:0')
tensor([ 0.9898, -0.4367,  0.2301, -0.5554, -0.2278], device='cuda:0')

tensor(1.2839, device='cuda:0', grad_fn=<MeanBackward0>)
tensor([[-0.0000, -0.0000, -0.0000, -0.1250, -0.1250],
        [-0.0000, -0.5000, -0.0000, -0.0000, -0.0000]], device='cuda:0')
tensor([[0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.],
        [0., 0.]], device='cuda:0')
tensor([[0., 0.]], device='cuda:0')
tensor([0.], device='cuda:0')
tensor([[0.],
        [0.],
        [0.],
        [0.],
        [0.]], device='cuda:0')
tensor([ 0.2957, -0.3807,  0.1376, -0.0131, -0.0396], device='cuda:0')