Repository: awslabs/sockeye

Issue:
hypotheses lengths off by one
It seems that the current beam search implementation (`_beam_search`) does not compute hypotheses length correctly. This can be seen when printing `sequences` and `lengths` just before returning from `_beam_search()`: the length is always one short of the actual non-zero values in sequences:

```
> sequences[0]
[    2    92  7741   907  2710   295    28     7  7741   907  2710     4
     7   383 10174  7324  9735   286   671   564  4085    71    22    98
 10671     4    13   782   485     8  6323  8665     4     8    23     8
  2976  1037  2448  3577    10   294   575  4664    13    34  3898   608
    43     5     3     0]
<NDArray 52 @cpu(0)>

> lengths[0]
[50.]
<NDArray 1 @cpu(0)>
```
`3` is the EOS id and the length of this hypothesis is clearly 51.

Let me try to explain the problem below. In practice, it doesn't matter since `_assemble_translation()` uses the length value to cut off the EOS symbol; the following EOS stripping logic in `_make_result()` is then redundant. But of course it shouldn't be and I think `_beam_search()` should return consistent data structures.

Here's the issue; I think this was introduced with length normalization only done for finished hypotheses:
1. `lengths` is initialized with 1s to account for the initial data in sequences, namely the BOS symbol
2. In each beam search iteration: as soon as we run topk on hypotheses scores, we effectively have expanded an hypothesis and its length increases by 1.
3. However, we are not using the updated length to potentially normalize the hypothesis if its finished. We check the newly acquired best_word_indices for EOS/PAD to compute a current `finished` status, and if so, **we normalize with the OLD length**.
4. This doesn't yet explain why we end beam_search with a length off by 1. This must somehow happen because the `length` array is only updated for hypotheses that are not finished and active. Since `finished` was just recomputed in step 3, the problem must lie in the `inactive` logic or the fact that right after we update `lengths`, we recompute `finished` once again and determine breaking out of the beam search loop based on that.

I am not 100% sure whats wrong here, probably the order of operations isn't entirely correct, but I also don't fully understand the whole `inactive` logic introduced. @mjpost could you comment on this?

At a minimum, we should make sure to normalize with correct length values. At best, we should make sure to return a `lengths` array from `_beam_search` that is consistent with `sequences`. With that, `_assemble_translation` would not silently cut off the last token (EOS), which by design should be handled later by `_make_result()`.



Discussion:
From looking at the code, I think the problem is:

- We normalize finished hypotheses, and then update the set of finished hypotheses.
- This prevents lengths from being updated.

I think it might be sufficient to:

- Return `all_finished` from the `NormalizeFinishedHypotheses` hybrid block (as is currently done, though without the extra variable assignment)
- In the calling block, save that to a new variable instead of overwriting `finished`. This value will then be used for pruning.

With that, the updating of lengths should work correctly because `finished` itself will not have been updated yet.

I should be able to look at this soon but thought I'd document this here in the meantime in case someone else wants to fix it.
Looking further, it seems we could eliminate the `UpdateLengthsAndFinished` hybrid block, and simply combine the two new blocks into a single block, done just before beam pruning.

Resolving patch:
diff --git a/sockeye/__init__.py b/sockeye/__init__.py
--- a/sockeye/__init__.py
+++ b/sockeye/__init__.py
@@ -11,4 +11,4 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-__version__ = '1.18.41'
+__version__ = '1.18.42'
diff --git a/sockeye/inference.py b/sockeye/inference.py
--- a/sockeye/inference.py
+++ b/sockeye/inference.py
@@ -896,42 +896,36 @@ def get(self, lengths: Union[mx.nd.NDArray, int, float]) -> Union[mx.nd.NDArray,
         return self.hybrid_forward(None, lengths)
 
 
-def _concat_translations(translations: List[Translation], start_id: int, stop_ids: Set[int],
+def _concat_translations(translations: List[Translation], stop_ids: Set[int],
                          length_penalty: LengthPenalty) -> Translation:
     """
     Combine translations through concatenation.
 
     :param translations: A list of translations (sequence starting with BOS symbol, attention_matrix), score and length.
-    :param start_id: The EOS symbol.
-    :param translations: The BOS symbols.
+    :param translations: The EOS symbols.
     :return: A concatenation if the translations with a score.
     """
     # Concatenation of all target ids without BOS and EOS
-    target_ids = [start_id]
+    target_ids = []
     attention_matrices = []
     beam_histories = []  # type: List[BeamHistory]
     for idx, translation in enumerate(translations):
-        assert translation.target_ids[0] == start_id
         if idx == len(translations) - 1:
-            target_ids.extend(translation.target_ids[1:])
-            attention_matrices.append(translation.attention_matrix[1:, :])
+            target_ids.extend(translation.target_ids)
+            attention_matrices.append(translation.attention_matrix)
         else:
             if translation.target_ids[-1] in stop_ids:
-                target_ids.extend(translation.target_ids[1:-1])
-                attention_matrices.append(translation.attention_matrix[1:-1, :])
+                target_ids.extend(translation.target_ids[:-1])
+                attention_matrices.append(translation.attention_matrix[:-1, :])
             else:
-                target_ids.extend(translation.target_ids[1:])
-                attention_matrices.append(translation.attention_matrix[1:, :])
+                target_ids.extend(translation.target_ids)
+                attention_matrices.append(translation.attention_matrix)
         beam_histories.extend(translation.beam_histories)
 
     # Combine attention matrices:
     attention_shapes = [attention_matrix.shape for attention_matrix in attention_matrices]
-    # Adding another row for the empty BOS alignment vector
-    bos_align_shape = np.asarray([1, 0])
-    attention_matrix_combined = np.zeros(np.sum(np.asarray(attention_shapes), axis=0) + bos_align_shape)
-
-    # We start at position 1 as position 0 is for the BOS, which is kept zero
-    pos_t, pos_s = 1, 0
+    attention_matrix_combined = np.zeros(np.sum(np.asarray(attention_shapes), axis=0))
+    pos_t, pos_s = 0, 0
     for attention_matrix, (len_t, len_s) in zip(attention_matrices, attention_shapes):
         attention_matrix_combined[pos_t:pos_t + len_t, pos_s:pos_s + len_s] = attention_matrix
         pos_t += len_t
@@ -1033,21 +1027,16 @@ def __init__(self,
         self._sort_by_index.initialize(ctx=self.context)
         self._sort_by_index.hybridize()
 
-        self._normalize_finished = NormalizeFinishedHypotheses(pad_id=C.PAD_ID,
-                                                               eos_id=self.vocab_target[C.EOS_SYMBOL],
-                                                               length_penalty_alpha=self.length_penalty.alpha,
-                                                               length_penalty_beta=self.length_penalty.beta)
-        self._normalize_finished.initialize(ctx=self.context)
-        self._normalize_finished.hybridize()
+        self._update_finished = NormalizeAndUpdateFinished(pad_id=C.PAD_ID,
+                                                           eos_id=self.vocab_target[C.EOS_SYMBOL],
+                                                           length_penalty_alpha=self.length_penalty.alpha,
+                                                           length_penalty_beta=self.length_penalty.beta)
+        self._update_finished.initialize(ctx=self.context)
+        self._update_finished.hybridize()
         self._prune_hyps = PruneHypotheses(threshold=self.beam_prune, beam_size=self.beam_size)
         self._prune_hyps.initialize(ctx=self.context)
         self._prune_hyps.hybridize()
 
-        self._update_lengths_and_finished = UpdateLengthsAndFinished(pad_id=C.PAD_ID,
-                                                                     eos_id=self.vocab_target[C.EOS_SYMBOL])
-        self._update_lengths_and_finished.initialize(ctx=self.context)
-        self._update_lengths_and_finished.hybridize()
-
         self.global_avoid_trie = None
         if avoid_list is not None:
             self.global_avoid_trie = constrained.AvoidTrie()
@@ -1250,9 +1239,8 @@ def _make_result(self,
         :param translation: The translation + attention and score.
         :return: TranslatorOutput.
         """
-        # remove special sentence start symbol (
) from the output:
-        target_ids = translation.target_ids[1:]
-        attention_matrix = translation.attention_matrix[1:, :]
+        target_ids = translation.target_ids
+        attention_matrix = translation.attention_matrix
 
         target_tokens = [self.vocab_target_inv[target_id] for target_id in target_ids]
 
@@ -1274,7 +1262,7 @@ def _concat_translations(self, translations: List[Translation]) -> Translation:
         :param translations: A list of translations (sequence, attention_matrix), score and length.
         :return: A concatenation if the translations with a score.
         """
-        return _concat_translations(translations, self.start_id, self.stop_ids, self.length_penalty)
+        return _concat_translations(translations, self.stop_ids, self.length_penalty)
 
     def _translate_nd(self,
                       source: mx.nd.NDArray,
@@ -1400,16 +1388,16 @@ def _beam_search(self,
 
         best_word_indices = mx.nd.full((self.batch_size * self.beam_size,), val=self.start_id, ctx=self.context,
                                        dtype='int32')
-        # growing sequences: (batch_size * beam_size, 1), pre-filled with 
 symbols on index 0
-        sequences = mx.nd.full((self.batch_size * self.beam_size, 1), val=self.start_id, ctx=self.context,
-                               dtype='int32')
+        # The growing sequences representing the complete history for each hypothesis: (batch_size * beam_size, 1)
+        # The first column represents 
, and is later removed when _beam_search() returns.
+        sequences = mx.nd.zeros((self.batch_size * self.beam_size, 1), ctx=self.context, dtype='int32')
 
         # Beam history
         beam_histories = None  # type: Optional[List[BeamHistory]]
         if self.store_beam:
             beam_histories = [defaultdict(list) for _ in range(self.batch_size)]
 
-        lengths = mx.nd.ones((self.batch_size * self.beam_size, 1), ctx=self.context)
+        lengths = mx.nd.zeros((self.batch_size * self.beam_size, 1), ctx=self.context)
         finished = mx.nd.zeros((self.batch_size * self.beam_size,), ctx=self.context, dtype='int32')
 
         # Extending max_output_lengths to shape (batch_size * beam_size,)
@@ -1533,10 +1521,11 @@ def _beam_search(self,
 
             # (5) Normalize the scores of newly finished hypotheses. Note that after this until the
             # next call to topk(), hypotheses may not be in sorted order.
-            finished, scores_accumulated = self._normalize_finished.forward(best_word_indices,
-                                                                            finished,
-                                                                            scores_accumulated,
-                                                                            lengths)
+            finished, scores_accumulated, lengths = self._update_finished.forward(best_word_indices,
+                                                                                  max_output_lengths,
+                                                                                  finished,
+                                                                                  scores_accumulated,
+                                                                                  lengths)
 
             # (6) Prune out low-probability hypotheses. Pruning works by setting entries `inactive`.
             if self.beam_prune > 0.0:
@@ -1546,18 +1535,12 @@ def _beam_search(self,
                                                                                            self.inf_array,
                                                                                            self.zeros_array)
 
-            # Update hypotheses lengths (unless finished or inactive) and compute finished hypotheses
-            finished, finished_or_inactive, lengths = self._update_lengths_and_finished.forward(finished,
-                                                                                                inactive,
-                                                                                                lengths,
-                                                                                                max_output_lengths,
-                                                                                                best_word_indices)
-
-            # (7) Reorder sequences and attentions according to best_hyp_indices
-            # and extend sequences with best_word_indices.
+            # (7) Reorder sequences and attentions according to best_hyp_indices...
             # pylint: disable=unsupported-assignment-operation
             sequences = mx.nd.take(sequences, best_hyp_indices)
             attentions = mx.nd.take(attentions, best_hyp_indices)
+
+            # ...and extend sequences with best_word_indices.
             sequences = mx.nd.concat(sequences, best_word_indices.reshape((-1, 1)), dim=1)
             attentions = mx.nd.concat(attentions, attention_scores.reshape((-1, 1, encoded_source_length)), dim=1)
 
@@ -1568,12 +1551,13 @@ def _beam_search(self,
 
             # (7b) optionally save beam history
             if self.store_beam:
+                finished_or_inactive = mx.nd.clip(data=finished + inactive, a_min=0, a_max=1)
                 unnormalized_scores = mx.nd.where(finished_or_inactive,
-                                                  scores_accumulated * self.length_penalty(lengths - 1),
+                                                  scores_accumulated * self.length_penalty(lengths),
                                                   scores_accumulated)
                 normalized_scores = mx.nd.where(finished_or_inactive,
                                                 scores_accumulated,
-                                                scores_accumulated / self.length_penalty(lengths - 1))
+                                                scores_accumulated / self.length_penalty(lengths))
                 for sent in range(self.batch_size):
                     rows = slice(sent * self.beam_size, (sent + 1) * self.beam_size)
 
@@ -1619,7 +1603,8 @@ def _beam_search(self,
         scores_accumulated = scores_accumulated.take(best_hyp_indices)
         constraints = [constraints[x] for x in best_hyp_indices.asnumpy()]
 
-        return sequences, attentions, scores_accumulated, lengths, constraints, beam_histories
+        # Remove the 
 column from both sequences and attentions
+        return sequences[:, 1:], attentions[:, 1:, :], scores_accumulated, lengths, constraints, beam_histories
 
     def _get_best_from_beam(self,
                             sequences: mx.nd.NDArray,
@@ -1754,7 +1739,7 @@ def hybrid_forward(self, F, indices, *args):
         return [F.take(arg, indices) for arg in args]
 
 
-class NormalizeFinishedHypotheses(mx.gluon.HybridBlock):
+class NormalizeAndUpdateFinished(mx.gluon.HybridBlock):
     """
     A HybridBlock for normalizing newly finished hypotheses scores with LengthPenalty.
     """
@@ -1769,31 +1754,18 @@ def __init__(self, pad_id: int,
         with self.name_scope():
             self.length_penalty = LengthPenalty(alpha=length_penalty_alpha, beta=length_penalty_beta)
 
-    def hybrid_forward(self, F, best_word_indices, finished, scores_accumulated, lengths):
+    def hybrid_forward(self, F, best_word_indices, max_output_lengths, finished, scores_accumulated, lengths):
         all_finished = ((best_word_indices == self.pad_id) + (best_word_indices == self.eos_id))
         newly_finished = all_finished - finished
         scores_accumulated = F.where(newly_finished,
                                      scores_accumulated / self.length_penalty(lengths),
                                      scores_accumulated)
-        finished = all_finished
-        return finished, scores_accumulated
-
 
-class UpdateLengthsAndFinished(mx.gluon.HybridBlock):
-    """
-    A HybridBlock that updates the lengths of unfinished and active hypotheses and recomputes finished hypotheses.
-    """
-
-    def __init__(self, pad_id: int, eos_id: int) -> None:
-        super().__init__()
-        self.pad_id = pad_id
-        self.eos_id = eos_id
+        # Update lengths of all items, except those that were already finished. This updates
+        # the lengths for inactive items, too, but that doesn't matter since they are ignored anyway.
+        lengths = lengths + F.cast(1 - F.expand_dims(finished, axis=1), dtype='float32')
 
-    def hybrid_forward(self, F, finished, inactive, lengths, max_output_lengths, best_word_indices):
-        finished_or_inactive = F.clip(data=finished + inactive, a_min=0, a_max=1)
-        # update lengths of hypotheses unless they are finished or inactive
-        lengths = lengths + F.cast(1 - F.expand_dims(finished_or_inactive, axis=1), dtype='float32')
-        # recompute finished. Hypotheses are finished if they are
+        # Now, recompute finished. Hypotheses are finished if they are
         # - extended with <pad>, or
         # - extended with <eos>, or
         # - at their maximum length.
@@ -1803,7 +1775,7 @@ def hybrid_forward(self, F, finished, inactive, lengths, max_output_lengths, bes
             (F.cast(F.reshape(lengths, shape=(-1,)), 'int32') >= max_output_lengths),
             a_min=0, a_max=1)
 
-        return finished, finished_or_inactive, lengths
+        return finished, scores_accumulated, lengths
 
 
 class UpdateScores(mx.gluon.HybridBlock):


Test patch:
diff --git a/test/unit/test_inference.py b/test/unit/test_inference.py
--- a/test/unit/test_inference.py
+++ b/test/unit/test_inference.py
@@ -73,7 +73,7 @@ def test_concat_translations():
     beam_history2 = {"id": [2]}
     beam_history3 = {"id": [3]}
     expected_beam_histories = [beam_history1, beam_history2, beam_history3]
-    expected_target_ids = [0, 1, 2, 8, 9, 3, 4, 5, -1]
+    expected_target_ids = [0, 1, 2, 0, 8, 9, 0, 3, 4, 5, -1]
     num_src = 7
 
     length_penalty = sockeye.inference.LengthPenalty()
@@ -87,8 +87,7 @@ def test_concat_translations():
                                                   [beam_history2]),
                     sockeye.inference.Translation([0, 3, 4, 5, -1], np.zeros((5, num_src)), 3.0 / length_penalty.get(5),
                                                   [beam_history3])]
-    combined = sockeye.inference._concat_translations(translations, start_id=_BOS, stop_ids={_EOS},
-                                                      length_penalty=length_penalty)
+    combined = sockeye.inference._concat_translations(translations, stop_ids={_EOS}, length_penalty=length_penalty)
 
     assert combined.target_ids == expected_target_ids
     assert combined.attention_matrix.shape == (len(expected_target_ids), len(translations) * num_src)