Training with masked matrix as input to softmax. Instead of setting inactive elements to zero, we need to set them to -inf.

masked_softmax.txt

th> x=torch.Tensor(2,3)
                                                                      [0.0000s]
th> x:uniform()
 0.2714  0.8744  0.5518
 0.3207  0.5189  0.8901
[torch.DoubleTensor of size 2x3]

                                                                      [0.0001s]
th> x[1][2] = -math.huge
                                                                      [0.0000s]
th> y=torch.LongTensor({1,3})
                                                                      [0.0000s]
th> ls=nn.LogSoftMax()
                                                                      [0.0001s]
th> c=nn.ClassNLLCriterion()
                                                                      [0.0001s]
th> c:forward(ls:forward(x), y)
0.82831845673287
                                                                      [0.0002s]
th> x:add(-0.01, ls:backward(x, c:backward(ls.output, y)))
  2.7426e-01         -inf   5.4894e-01
  3.1948e-01   5.1737e-01   8.9287e-01
[torch.DoubleTensor of size 2x3]

                                                                      [0.0002s]
th> c:forward(ls:forward(x), y)
0.82553313455311
                                                                      [0.0002s]
th> x:add(-0.01, ls:backward(x, c:backward(ls.output, y)))
  2.7710e-01         -inf   5.4610e-01
  3.1823e-01   5.1584e-01   8.9565e-01
[torch.DoubleTensor of size 2x3]

                                                                      [0.0002s]
th> c:forward(ls:forward(x), y)
0.82276008636199
                                                                      [0.0002s]
th> x:add(-0.01, ls:backward(x, c:backward(ls.output, y)))
  2.7993e-01         -inf   5.4326e-01
  3.1698e-01   5.1432e-01   8.9843e-01
[torch.DoubleTensor of size 2x3]

                                                                      [0.0002s]
th> c:forward(ls:forward(x), y)
0.81999926513474
                                                                      [0.0001s]
th> x:add(-0.01, ls:backward(x, c:backward(ls.output, y)))
  2.8276e-01         -inf   5.4044e-01
  3.1573e-01   5.1280e-01   9.0119e-01
[torch.DoubleTensor of size 2x3]

                                                                      [0.0002s]
th> c:forward(ls:forward(x), y)
0.81725062381229
                                                                      [0.0001s]
th> x:add(-0.01, ls:backward(x, c:backward(ls.output, y)))
  2.8558e-01         -inf   5.3762e-01
  3.1449e-01   5.1128e-01   9.0396e-01
[torch.DoubleTensor of size 2x3]

                                                                      [0.0002s]
th> c:forward(ls:forward(x), y)
0.81451411530494
                                                                      [0.0001s]
th> x:add(-0.01, ls:backward(x, c:backward(ls.output, y)))
  2.8839e-01         -inf   5.3480e-01
  3.1324e-01   5.0977e-01   9.0671e-01
[torch.DoubleTensor of size 2x3]

                                                                      [0.0002s]