Repository: espnet/espnet

Issue:
RNNT modified_adaptive_expansion_search is too slow.
Hi, I try to use `modified_adaptive_expansion_search` with beam_size 5 to recognize one wav, which costs 52.9s, while `ALSD` costs 0.87S. why the speed is so slow. 
RNNT modified_adaptive_expansion_search is too slow.
Hi, I try to use `modified_adaptive_expansion_search` with beam_size 5 to recognize one wav, which costs 52.9s, while `ALSD` costs 0.87S. why the speed is so slow. 


Discussion:
Hi,

Can you give information? recipe, architecture, units, etc?

Nevertheless, I'm assuming you are modeling BPE or use a big vocab size? If so, the issue is known for ESPnet1 and due to `select_k_expansions(...)`, see [here](https://github.com/espnet/espnet/pull/4032#issuecomment-1073682573)

> mAES: Fixed select_top_k_candidates loop taking a long time with a big vocab size. It also improves decoding time in general.

I'll fix it in ESPnet1 when I have time, for the moment please either use the changes from PR #4032 or use another time-synchronous algorithm. 
Thank you for your reply, The dataset is AISHELL, vocab_size = 4442, encoder consist of 11 blocks. And I tried your PR, find the new implement cost 0.32s, which is faster than `ALSD`. but another problem is that the hyps in final beam are all same(beam_width=5). Can you take a look ? 
> And I tried your PR, find the new implement cost 0.32s, which is faster than ALSD. but another problem is that the hyps in final beam are all same(beam_width=5). Can you take a look ?

Oh, thank you for asking that! The original algorithm (in ESPnet1) doesn't include a duplication check but I wanted to try adding the merging procedure (from ALSD) in #4032, applied at each timestep. Could you please test replacing the following lines:

```python
  kept_hyps = sorted(list_b, key=lambda x: x.score, reverse=True)[: self.beam_size]
  #... and
  kept_hyps = sorted(list_b + list_exp, key=lambda x: x.score, reverse=True)[: self.beam_size]
```

with:

```python
  kept_hyps = self.recombine_hyps(sorted(list_b, key=lambda x: x.score, reverse=True))[: self.beam_size]
  # ... and
  kept_hyps = self.recombine_hyps(sorted(list_b + list_exp, key=lambda x: x.score, reverse=True))[: self.beam_size]
```

Note that depending on `expansion_gamma` and `expansion_beta` values, the pruning may be aggressive. Thus, you can end-up with 1) less final hypothesis than specified and 2) the final hypothesis will be almost identical.
I will try it today.<br>
```
kept_hyps = sorted(self.recombine_hyps(list_b), key=lambda x: x.score, reverse=True)[: self.beam_size]
# ... and
kept_hyps = sorted(self.recombine_hyps(list_b + list_exp), key=lambda x: x.score, reverse=True)[: `self.beam_size
```
Should `recombine_hyps` first, and then `sorted` ? 
  
It shouldn't make a difference in practice but yes, it should be `recombine_hyps` + `sorted` instead of `sorted` + `recombine_hyps` (pruning should be applied at the end anyway). Will change that in #4032.
Hi,

Can you give information? recipe, architecture, units, etc?

Nevertheless, I'm assuming you are modeling BPE or use a big vocab size? If so, the issue is known for ESPnet1 and due to `select_k_expansions(...)`, see [here](https://github.com/espnet/espnet/pull/4032#issuecomment-1073682573)

> mAES: Fixed select_top_k_candidates loop taking a long time with a big vocab size. It also improves decoding time in general.

I'll fix it in ESPnet1 when I have time, for the moment please either use the changes from PR #4032 or use another time-synchronous algorithm. 
Thank you for your reply, The dataset is AISHELL, vocab_size = 4442, encoder consist of 11 blocks. And I tried your PR, find the new implement cost 0.32s, which is faster than `ALSD`. but another problem is that the hyps in final beam are all same(beam_width=5). Can you take a look ? 
> And I tried your PR, find the new implement cost 0.32s, which is faster than ALSD. but another problem is that the hyps in final beam are all same(beam_width=5). Can you take a look ?

Oh, thank you for asking that! The original algorithm (in ESPnet1) doesn't include a duplication check but I wanted to try adding the merging procedure (from ALSD) in #4032, applied at each timestep. Could you please test replacing the following lines:

```python
  kept_hyps = sorted(list_b, key=lambda x: x.score, reverse=True)[: self.beam_size]
  #... and
  kept_hyps = sorted(list_b + list_exp, key=lambda x: x.score, reverse=True)[: self.beam_size]
```

with:

```python
  kept_hyps = self.recombine_hyps(sorted(list_b, key=lambda x: x.score, reverse=True))[: self.beam_size]
  # ... and
  kept_hyps = self.recombine_hyps(sorted(list_b + list_exp, key=lambda x: x.score, reverse=True))[: self.beam_size]
```

Note that depending on `expansion_gamma` and `expansion_beta` values, the pruning may be aggressive. Thus, you can end-up with 1) less final hypothesis than specified and 2) the final hypothesis will be almost identical.
I will try it today.<br>
```
kept_hyps = sorted(self.recombine_hyps(list_b), key=lambda x: x.score, reverse=True)[: self.beam_size]
# ... and
kept_hyps = sorted(self.recombine_hyps(list_b + list_exp), key=lambda x: x.score, reverse=True)[: `self.beam_size
```
Should `recombine_hyps` first, and then `sorted` ? 
  
It shouldn't make a difference in practice but yes, it should be `recombine_hyps` + `sorted` instead of `sorted` + `recombine_hyps` (pruning should be applied at the end anyway). Will change that in #4032.

Resolving patch:
diff --git a/espnet/nets/beam_search_transducer.py b/espnet/nets/beam_search_transducer.py
--- a/espnet/nets/beam_search_transducer.py
+++ b/espnet/nets/beam_search_transducer.py
@@ -95,7 +95,13 @@ def __init__(
             self.nstep = nstep if nstep > 1 else 2
             self.prefix_alpha = prefix_alpha
             self.expansion_gamma = expansion_gamma
-            self.expansion_beta = expansion_beta
+
+            assert self.vocab_size >= beam_size + expansion_beta, (
+                "beam_size (%d) + expansion_beta (%d) "
+                "should be smaller or equal to vocabulary size (%d)."
+                % (beam_size, expansion_beta, self.vocab_size)
+            )
+            self.max_candidates = beam_size + expansion_beta
 
             self.search_algorithm = self.modified_adaptive_expansion_search
         else:
@@ -779,16 +785,22 @@ def modified_adaptive_expansion_search(
             beam_enc_out = enc_out_t.unsqueeze(0)
 
             list_b = []
+            duplication_check = [hyp.yseq for hyp in hyps]
+
             for n in range(self.nstep):
                 beam_dec_out = torch.stack([h.dec_out[-1] for h in hyps])
 
-                beam_logp = torch.log_softmax(
+                beam_logp, beam_idx = torch.log_softmax(
                     self.joint_network(beam_enc_out, beam_dec_out)
                     / self.softmax_temperature,
                     dim=-1,
-                )
+                ).topk(self.max_candidates, dim=-1)
+
                 k_expansions = select_k_expansions(
-                    hyps, beam_logp, beam, self.expansion_gamma, self.expansion_beta
+                    hyps,
+                    beam_idx,
+                    beam_logp,
+                    self.expansion_gamma,
                 )
 
                 list_exp = []
@@ -806,14 +818,15 @@ def modified_adaptive_expansion_search(
                         if k == 0:
                             list_b.append(new_hyp)
                         else:
-                            new_hyp.yseq.append(int(k))
+                            if new_hyp.yseq + [int(k)] not in duplication_check:
+                                new_hyp.yseq.append(int(k))
 
-                            if self.use_lm:
-                                new_hyp.score += self.lm_weight * float(
-                                    hyp.lm_scores[k]
-                                )
+                                if self.use_lm:
+                                    new_hyp.score += self.lm_weight * float(
+                                        hyp.lm_scores[k]
+                                    )
 
-                            list_exp.append(new_hyp)
+                                list_exp.append(new_hyp)
 
                 if not list_exp:
                     kept_hyps = sorted(list_b, key=lambda x: x.score, reverse=True)[
diff --git a/espnet/nets/pytorch_backend/transducer/utils.py b/espnet/nets/pytorch_backend/transducer/utils.py
--- a/espnet/nets/pytorch_backend/transducer/utils.py
+++ b/espnet/nets/pytorch_backend/transducer/utils.py
@@ -136,10 +136,9 @@ def subtract(
 
 def select_k_expansions(
     hyps: List[ExtendedHypothesis],
-    logps: torch.Tensor,
-    beam_size: int,
+    topk_idxs: torch.Tensor,
+    topk_logps: torch.Tensor,
     gamma: float,
-    beta: float,
 ) -> List[ExtendedHypothesis]:
     """Return K hypotheses candidates for expansion from a list of hypothesis.
 
@@ -148,10 +147,9 @@ def select_k_expansions(
 
     Args:
         hyps: Hypotheses.
-        beam_logp: Log-probabilities for hypotheses expansions.
-        beam_size: Beam size.
+        topk_idxs: Indices of candidates hypothesis.
+        topk_logps: Log-probabilities for hypotheses expansions.
         gamma: Allowed logp difference for prune-by-value method.
-        beta: Number of additional candidates to store.
 
     Return:
         k_expansions: Best K expansion hypotheses candidates.
@@ -160,7 +158,9 @@ def select_k_expansions(
     k_expansions = []
 
     for i, hyp in enumerate(hyps):
-        hyp_i = [(int(k), hyp.score + float(logp)) for k, logp in enumerate(logps[i])]
+        hyp_i = [
+            (int(k), hyp.score + float(v)) for k, v in zip(topk_idxs[i], topk_logps[i])
+        ]
         k_best_exp = max(hyp_i, key=lambda x: x[1])[1]
 
         k_expansions.append(
@@ -168,7 +168,7 @@ def select_k_expansions(
                 filter(lambda x: (k_best_exp - gamma) <= x[1], hyp_i),
                 key=lambda x: x[1],
                 reverse=True,
-            )[: beam_size + beta]
+            )
         )
 
     return k_expansions
diff --git a/espnet2/asr/transducer/beam_search_transducer.py b/espnet2/asr/transducer/beam_search_transducer.py
--- a/espnet2/asr/transducer/beam_search_transducer.py
+++ b/espnet2/asr/transducer/beam_search_transducer.py
@@ -59,7 +59,7 @@ def __init__(
         expansion_beta: int = 2,
         score_norm: bool = True,
         nbest: int = 1,
-        token_list: List[str] = None,
+        token_list: Optional[List[str]] = None,
     ):
         """Initialize Transducer search module.
 
@@ -100,18 +100,21 @@ def __init__(
         elif search_type == "tsd":
             if isinstance(lm, TransformerLM):
                 raise NotImplementedError
+
             self.max_sym_exp = max_sym_exp
 
             self.search_algorithm = self.time_sync_decoding
         elif search_type == "alsd":
             if isinstance(lm, TransformerLM):
                 raise NotImplementedError
+
             self.u_max = u_max
 
             self.search_algorithm = self.align_length_sync_decoding
         elif search_type == "nsc":
             if isinstance(lm, TransformerLM):
                 raise NotImplementedError
+
             self.nstep = nstep
             self.prefix_alpha = prefix_alpha
 
@@ -119,10 +122,17 @@ def __init__(
         elif search_type == "maes":
             if isinstance(lm, TransformerLM):
                 raise NotImplementedError
+
             self.nstep = nstep if nstep > 1 else 2
             self.prefix_alpha = prefix_alpha
             self.expansion_gamma = expansion_gamma
-            self.expansion_beta = expansion_beta
+
+            assert self.vocab_size >= beam_size + expansion_beta, (
+                "beam_size (%d) + expansion_beta (%d) "
+                "should be smaller or equal to vocabulary size (%d)."
+                % (beam_size, expansion_beta, self.vocab_size)
+            )
+            self.max_candidates = beam_size + expansion_beta
 
             self.search_algorithm = self.modified_adaptive_expansion_search
         else:
@@ -279,6 +289,7 @@ def default_beam_search(self, enc_out: torch.Tensor) -> List[Hypothesis]:
                         ]
                     )
                 )
+
             while True:
                 max_hyp = max(hyps, key=lambda x: x.score)
                 hyps.remove(max_hyp)
@@ -774,15 +785,21 @@ def modified_adaptive_expansion_search(
             beam_enc_out = enc_out_t.unsqueeze(0)
 
             list_b = []
+            duplication_check = [hyp.yseq for hyp in hyps]
+
             for n in range(self.nstep):
                 beam_dec_out = torch.stack([h.dec_out[-1] for h in hyps])
 
-                beam_logp = torch.log_softmax(
+                beam_logp, beam_idx = torch.log_softmax(
                     self.joint_network(beam_enc_out, beam_dec_out),
                     dim=-1,
-                )
+                ).topk(self.max_candidates, dim=-1)
+
                 k_expansions = select_k_expansions(
-                    hyps, beam_logp, beam, self.expansion_gamma, self.expansion_beta
+                    hyps,
+                    beam_idx,
+                    beam_logp,
+                    self.expansion_gamma,
                 )
 
                 list_exp = []
@@ -800,14 +817,15 @@ def modified_adaptive_expansion_search(
                         if k == 0:
                             list_b.append(new_hyp)
                         else:
-                            new_hyp.yseq.append(int(k))
+                            if new_hyp.yseq + [int(k)] not in duplication_check:
+                                new_hyp.yseq.append(int(k))
 
-                            if self.use_lm:
-                                new_hyp.score += self.lm_weight * float(
-                                    hyp.lm_scores[k]
-                                )
+                                if self.use_lm:
+                                    new_hyp.score += self.lm_weight * float(
+                                        hyp.lm_scores[k]
+                                    )
 
-                            list_exp.append(new_hyp)
+                                list_exp.append(new_hyp)
 
                 if not list_exp:
                     kept_hyps = sorted(list_b, key=lambda x: x.score, reverse=True)[


Test patch: