Repository: facebookresearch/ParlAI

Issue:
MemNN padding issue
I'm seeing different results with MemNN while varying batch size.  The problem seems to relate to attention and padding.

If I have a memory that gets padded (due to another training sample in its batch being longer), we add 0's to the end of the memory so that their dimensions match.  When we [do attention](https://github.com/facebookresearch/ParlAI/blob/master/parlai/agents/memnn/modules.py#L203-L204) we don't treat these paddings any differently which will lead to different results in `probs` after applying the softmax.  For example:

```Python
>>> F.softmax(torch.Tensor([1, 4, 5, 3]))
tensor([0.0120, 0.2418, 0.6572, 0.0889])
>>> F.softmax(torch.Tensor([1, 4, 5, 3, 0, 0, 0]))
tensor([0.0119, 0.2386, 0.6486, 0.0878, 0.0044, 0.0044, 0.0044])
>>>
```

This causes the results to diverge at each hop.  [Fairseq gets around this](https://github.com/pytorch/fairseq/blob/master/fairseq/modules/multihead_attention.py#L186-L189) by setting padded values to `-inf`

```Python
>>> F.softmax(torch.Tensor([1, 4, 5, 3, float('-inf'), float('-inf'), float('-inf')]))
tensor([0.0120, 0.2418, 0.6572, 0.0889, 0.0000, 0.0000, 0.0000])
```




Discussion:
Another thing to keep in mind is that [we add time features to padded memories](https://github.com/facebookresearch/ParlAI/blob/master/parlai/agents/memnn/memnn.py#L191).  We'll probably want to fix this as well.  I'm happy to submit a PR, but maybe we should discuss further.
Historically MemNN is only trained using hogwild using batch size of 1. I'm not surprised there are padding issues. A PR would be greatly appreciated (and a test!)

Resolving patch:
diff --git a/parlai/agents/memnn/memnn.py b/parlai/agents/memnn/memnn.py
--- a/parlai/agents/memnn/memnn.py
+++ b/parlai/agents/memnn/memnn.py
@@ -52,7 +52,8 @@ def add_cmdline_args(argparser):
 
     @staticmethod
     def model_version():
-        """Return current version of this model, counting up from 0.
+        """
+        Return current version of this model, counting up from 0.
 
         Models may not be backwards-compatible with older versions.
         Version 1 split from version 0 on Sep 7, 2018.
@@ -87,7 +88,11 @@ def build_model(self):
 
     def score_candidates(self, batch, cand_vecs, cand_encs=None):
         mems = self._build_mems(batch.memory_vecs)
-        scores = self.model(batch.text_vec, mems, cand_vecs)
+        # Check for rows that have no non-null tokens
+        pad_mask = None
+        if mems is not None:
+            pad_mask = (mems != self.NULL_IDX).sum(dim=-1) == 0
+        scores = self.model(batch.text_vec, mems, cand_vecs, pad_mask)
         return scores
 
     @lru_cache(maxsize=None)  # bounded by opt['memsize'], cache string concats
@@ -153,14 +158,17 @@ def _set_text_vec(self, obs, history, truncate):
         return obs
 
     def _build_mems(self, mems):
-        """Build memory tensors.
+        """
+        Build memory tensors.
 
         During building, will add time features to the memories if enabled.
 
-        :param: list of length batchsize containing inner lists of 1D tensors
-                containing the individual memories for each row in the batch.
+        :param mems:
+            list of length batchsize containing inner lists of 1D tensors
+            containing the individual memories for each row in the batch.
 
-        :returns: 3d padded tensor of memories (bsz x num_mems x seqlen)
+        :returns:
+            3d padded tensor of memories (bsz x num_mems x seqlen)
         """
         if mems is None:
             return None
@@ -186,20 +194,11 @@ def _build_mems(self, mems):
         padded = torch.LongTensor(bsz, num_mems, seqlen).fill_(0)
 
         for i, mem in enumerate(mems):
-            # tf_offset = len(mem) - 1
+            tf_offset = len(mem) - 1
             for j, m in enumerate(mem):
                 padded[i, j, :len(m)] = m
-                # if self.use_time_features:
-                #     padded[i, j, -1] = self.dict[self._time_feature(tf_offset - j)]
-
-        # NOTE: currently below we are adding tf's to every memory,
-        # including emtpy ones. above commented-out code adds only to filled
-        # ones but is significantly slower to run.
-        if self.use_time_features:
-            nm = num_mems - 1
-            for i in range(num_mems):
-                # put lowest time feature in most recent memory
-                padded[:, nm - i, -1] = self.dict[self._time_feature(i)]
+                if self.use_time_features:
+                    padded[i, j, -1] = self.dict[self._time_feature(tf_offset - j)]
 
         if self.use_cuda:
             padded = padded.cuda()
diff --git a/parlai/agents/memnn/modules.py b/parlai/agents/memnn/modules.py
--- a/parlai/agents/memnn/modules.py
+++ b/parlai/agents/memnn/modules.py
@@ -6,7 +6,7 @@
 
 import torch
 import torch.nn as nn
-
+from parlai.core.utils import neginf
 from functools import lru_cache
 
 
@@ -61,19 +61,28 @@ def _score(self, output, cands):
             raise RuntimeError('Unexpected candidate dimensions {}'
                                ''.format(cands.dim()))
 
-    def forward(self, xs, mems, cands=None):
-        """One forward step.
+    def forward(self, xs, mems, cands=None, pad_mask=None):
+        """
+        One forward step.
+
+        :param xs:
+            (bsz x seqlen) LongTensor queries to the model
+
+        :param mems:
+            (bsz x num_mems x seqlen) LongTensor memories
 
-        :param xs:    (bsz x seqlen) LongTensor queries to the model
-        :param mems:  (bsz x num_mems x seqlen) LongTensor memories
-        :param cands: (num_cands x seqlen) or (bsz x num_cands x seqlen)
-                      LongTensor with candidates to rank
+        :param cands:
+            (num_cands x seqlen) or (bsz x num_cands x seqlen)
+            LongTensor with candidates to rank
+        :param pad_mask:
+            (bsz x num_mems) optional mask indicating which tokens
+            correspond to padding
 
-        :returns: scores
+        :returns:
             scores contains the model's predicted scores.
-                if cand_params is None, the candidates are the vocabulary;
-                otherwise, these scores are over the candidates provided.
-                (bsz x num_cands)
+            if cand_params is None, the candidates are the vocabulary;
+            otherwise, these scores are over the candidates provided.
+            (bsz x num_cands)
         """
         state = self.query_lt(xs)
         if mems is not None:
@@ -82,7 +91,8 @@ def forward(self, xs, mems, cands=None):
             out_memory_embs = self.out_memory_lt(mems)
 
             for _ in range(self.hops):
-                state = self.memory_hop(state, in_memory_embs, out_memory_embs)
+                state = self.memory_hop(state, in_memory_embs, out_memory_embs,
+                                        pad_mask)
 
         if cands is not None:
             # embed candidates
@@ -96,18 +106,22 @@ def forward(self, xs, mems, cands=None):
 
 
 class Embed(nn.Embedding):
-    """Embed sequences for MemNN model.
+    """
+    Embed sequences for MemNN model.
 
     Applies Position Encoding if enabled and currently applies BOW sum.
     """
 
     def __init__(self, *args, position_encoding=False, reduction='mean',
                  **kwargs):
-        """Initialize custom Embedding layer.
+        """
+        Initialize custom Embedding layer.
+
+        :param position_encoding:
+            apply positional encoding transformation on input sequences
 
-        :param position_encoding: apply positional encoding transformation
-                                  on input sequences
-        :param reduction:         reduction strategy to sequences, default 'mean'
+        :param reduction:
+            reduction strategy to sequences, default 'mean'
         """
         self.position_encoding = position_encoding
         self.reduction = reduction
@@ -120,18 +134,22 @@ def _reduce(self, embs, input):
         elif self.reduction == 'mean':
             # this is more fair than mean(-2) since mean includes null tokens
             sum = embs.sum(-2)
-            lens = input.ne(0).sum(-1).unsqueeze(-1).float()
+            lens = input.ne(self.padding_idx).sum(-1).unsqueeze(-1).float()\
+                .clamp_(min=1)
             return sum / lens
         else:
             raise RuntimeError(
                 'reduction method {} not supported'.format(self.reduction))
 
     def forward(self, input):
-        """Return BOW embedding with PE reweighting if enabled.
+        """
+        Return BOW embedding with PE reweighting if enabled.
 
-        :param input: (bsz x seqlen) LongTensor
+        :param input:
+            (bsz x seqlen) LongTensor
 
-        :returns: (bsz x esz) FloatTensor
+        :returns:
+            (bsz x esz) FloatTensor
         """
         embs = super().forward(input)
         if self.position_encoding:
@@ -150,15 +168,20 @@ def forward(self, input):
     @staticmethod
     @lru_cache(maxsize=128)
     def position_matrix(J, d, use_cuda):
-        """Build matrix of position encoding coeffiencents.
+        """
+        Build matrix of position encoding coeffiencents.
 
         See https://papers.nips.cc/paper/5846-end-to-end-memory-networks,
         section 4.1 Model Details: Sentence Representation.
 
-        :param J: number of words in the sequence
-        :param d: dimension of the embedding
+        :param J:
+            number of words in the sequence
 
-        :returns: Position Encoding matrix
+        :param d:
+            dimension of the embedding
+
+        :returns:
+            Position Encoding matrix
         """
         m = torch.Tensor(J, d)
         for k in range(1, d + 1):
@@ -188,19 +211,30 @@ def __init__(self, embedding_size, rotate=True):
             self.rotate = lambda x: x
         self.softmax = nn.Softmax(dim=1)
 
-    def forward(self, query_embs, in_mem_embs, out_mem_embs):
-        """Compute MemNN Hop step.
+    def forward(self, query_embs, in_mem_embs, out_mem_embs, pad_mask):
+        """
+        Compute MemNN Hop step.
+
+        :param query_embs:
+            (bsz x esz) embedding of queries
+
+        :param in_mem_embs:
+            bsz list of (num_mems x esz) embedding of memories for activation
+
+        :param out_mem_embs:
+            bsz list of (num_mems x esz) embedding of memories for outputs
 
-        :param query_embs:   (bsz x esz) embedding of queries
-        :param in_mem_embs:  bsz list of (num_mems x esz) embedding of memories
-                             for activation
-        :param out_mem_embs: bsz list of (num_mems x esz) embedding of memories
-                             for outputs
+        :param pad_mask
+            (bsz x num_mems) optional mask indicating which tokens correspond to
+            padding
 
-        :returns: (bsz x esz) output state
+        :returns:
+            (bsz x esz) output state
         """
         # rotate query embeddings
         attn = torch.bmm(query_embs.unsqueeze(1), in_mem_embs).squeeze(1)
+        if pad_mask is not None:
+            attn[pad_mask] = neginf(attn.dtype)
         probs = self.softmax(attn)
         memory_output = torch.bmm(probs.unsqueeze(1), out_mem_embs).squeeze(1)
         output = memory_output + self.rotate(query_embs)


Test patch: