Notes

몇 년 전, Alex Graves가 길이가 다른 input/output sequences 를 잘 mapping 할 수 있는 RNN-T 모델을 Sequence Transduction with RNN [1] 논문에서 제안했다. 여러 가지 sequence transduction 문제 중 ASR 문제에 집중을 했는데, 구글이 작년에 이 모델로 모바일 상에서 steamable 음성인식 성과를 소개해서인지, 최근들어 이 모델에 대한 관심이 급격하게 높아졌다. ASR 문제에서 input sequence는 audio signal, output sequence는 텍스트이다. RNN-T에 대한 설명은 생략하고, 이 글의 주제인 beam search for RNN-T 알고리즘에 대해서 메모를 남겨보고자 한다. 이 알고리즘은 [1] 논문에서 처음 제안되었다. pseudo code는 다음과 같다.

대표사진 삭제

사진 설명을 입력하세요.

RNN-T 모델을 위한 beam search를 위해 우선 time t에서 hypothesis h 는 frame embedding a_t 와 text embedding t_h 를 가진다고 가정한다. 이 두 embedding들은 결합되고 Joiner를 태워서 y_t,h output units(eg, 256) 에 대한 확률을 생성한다. while 문 안에 있는 for문에서의 y는 한 시점의 output units 이다.

RNN-T의 output units 은 우리가 모델링하고자 하는 output character들 외에 특별한 symbol 인 blank (∅) 가 존재한다. 이 blank symbol 은 디코딩 진행시, lattice 상에서 t 시점에 있을 때, 다음 time frame인 t+1 으로 진행할 것인지 혹은 t 시점에서의 output unit 과 같은 문자를 emit할 것인지 결정하기 위해 존재한다.

Beam search는 두 가지 hypothesis set A, B 를 사용하여 수행이 된다. 이것은 t + 1로 이동되는 최소 W (beam size) hypothesis가 t에서 여전히 생성 될 수있는 것보다 높은 확률을 갖도록하기 위함이다.

알고리즘을 보면 크게 세 부분으로 나눌 수 있다.

1. 첫번째는 특정시점까지 찾아진 hypothesis set A에 포함된 모든 hypothesis에 대한 probability를 그 hypothesis 의 적절한 prefix의 set을 사용하여 계산해 내는, 첫번째 for-loop에 대한 부분이다. 모든 time-step마다 그 시점까지 찾아진 hypothesis set B(=A) 에 포함된 모든 hypothesis에 대한 확률을 계산해낸다.

2. 두번째는 Beam search 를 진행하는 중에, hypothesis set A에서 최고 확률이 높은 hypothesis를 선택하고 blank 또는 non-blank symbol 로 확장하는 부분이다. while 를 돌며 Joiner(softmax)를 태워서 나온 y_t,h output units(eg, 256) 에서 모든 dimension output(probability)에 대해서, 기존 score(prob)에 곱한 score(prob)를 업데이트 하고, 해당 chracter(정확히는 index)를 append 시킨 새로운 hypothesis를 hypothesis set A 혹은 B에 append 한다. 이 때, blank에 해당하는 output dimension에 대해서는 score 계산 후, hypothesis set B에 append 하고, 다른 dimension에 해당하는 값(다른 캐릭터) 들에 대해서는 socre 계산 후, hypothesis set A에 append한다. While 문은 beam width 만큼 hypothesis가 hypothesis set B에 존재할 때까지 진행된다. 다시 말하면, blank가 마지막에 사용되는 expansion hypothesis는 를 hypo set B에 삽입된다. 한편, blank가 사용되지 않는 다른 모든 symbols을 사용하는 expansion hypothesis들은 hypothesis set A에 다시 삽입되어 t에서 확장 된 hypothesis set A가 계속해서 재 정의 된다.

3. return 할 때, hypothesis set B에 속한 W 개 미만의 hypothesis 중에 "각각의 hypotesis의 확률 값을 chracter들의 숫자로 나눈 최종 score"을 비교하여 가장 높은 score 값이 되는 hypothesis를 return 한다.

** 참고로 hypothesis set A, B에 저장 할 때는 {score(probability), character sequence(hypothesis)} 가 사전 형태로 함께 저장된다는 것을 함께 생각해야, 헷갈리지 않는다.

espnet에 구현되어 있는 beam search 코드[2]이다. 참고하면 더 이해하기 편하다. pseudo code와는 다르게 코드 상에선 blank output 에 대해서 해당 hypothesis를 hypothesis set B 에 append 하는 부분을 for 문 안으로 넣어놓았다.

def recognize_beam(self, h, recog_args, rnnlm=None):
    """Beam search implementation for transformer-transducer.
    Args:
        h (torch.Tensor): encoder hidden state sequences (maxlen_in, Henc)
        recog_args (Namespace): argument Namespace containing options
        rnnlm (torch.nn.Module): language model module
    Returns:
        nbest_hyps (list of dicts): n-best decoding results
    """
    beam = recog_args.beam_size
    k_range = min(beam, self.odim)
    nbest = recog_args.nbest
    normscore = recog_args.score_norm_transducer

    B_hyps = [{'score': 0.0, 'yseq': [self.blank], 'cache': None}]

    for i, hi in enumerate(h):
        A_hyps = B_hyps
        B_hyps = []

        while True:
            new_hyp = max(A_hyps, key=lambda x: x['score'])
            A_hyps.remove(new_hyp)

            ys = to_device(self, torch.tensor(new_hyp['yseq']).unsqueeze(0))
            ys_mask = to_device(self, subsequent_mask(len(new_hyp['yseq'])).unsqueeze(0))
            y, c = self.forward_one_step(ys, ys_mask, new_hyp['cache'])

            ytu = torch.log_softmax(self.joint(hi, y[0]), dim=0)

            for k in six.moves.range(self.odim):
                beam_hyp = {'score': new_hyp['score'] + float(ytu[k]),
                            'yseq': new_hyp['yseq'][:],
                            'cache': new_hyp['cache']}

                if k == self.blank:
                    B_hyps.append(beam_hyp)
                else:
                    beam_hyp['yseq'].append(int(k))
                    beam_hyp['cache'] = c

                    A_hyps.append(beam_hyp)

            if len(B_hyps) >= k_range: // beam_size (W)
                break

    if normscore:
        nbest_hyps = sorted(
            B_hyps, key=lambda x: x['score'] / len(x['yseq']), reverse=True)[:nbest]
    else:
        nbest_hyps = sorted(
            B_hyps, key=lambda x: x['score'], reverse=True)[:nbest]

    return nbest_hyps

[1] Alex Graves, "Sequence Transduction with Recurrent Neural Networks", 2012

[2] ESPnet; https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/transducer/transformer_decoder.py