Prefix Beam Search

发表于 更新于

一个急于实现的意外收获，以及再次遇到动态规划和深度学习的结合。

prefix beam search在语音识的解码阶段有着重要的作用。并且提到prefix beam search不得不提一下CTC的主要思路。

强大的深度学习进军语音识别领域，企图取代传统的GMM-HMM + ngram LM架构，神经网络模型的输入一般是有部分重叠的声音片段（长度为几十微秒左右）序列，然后模型直接输出预测的音素序列。不过输出音素序列一般较长，并且有较多的重复，就比如一个单词hello的输出音素可能是[h,h,e,e,e,l,l,l,o,o,o]（这里简单起见就用字母直接代替了音素）。

那么问题来了，怎么使用一种简单的方式对齐模型输出的音素和正确的音素呢？其实这个问题很容易让人联想到动态规划，输出串和目标串之间的编辑距离嘛（大致思路）。不过还有一个问题，训练的时候有目标串给出对齐方式，但是预测的时候怎么办呢？没有了参考的对齐方式，就像上面那个例子，hello中连续的两个l怎么区分，是预测输出helo还是hello？。CTC的解决方案是通过引入空白符号blank，即下划线_来分隔连续且相同的字符，[h,h,_,e,e,e,l,_,l,l,o,o,o]就很容易通过“消消乐”的手段看出来是hello了嘛。

接下来要讲的重点是如何从模型的输出中解码出最有可能的序列。

Greedy Search

一种一拍脑袋就能想出来的解码方式是greedy search，即直接挑选每帧中概率最大的字符。那么如果给定seq_len × n_class的输入，输出的字符串长度一般就是seq_len了。这么做显然忽视了当前帧和前后其他帧之间的关联，但是如果模型使用了bi-LSTM这种能同时考虑上下文的语言模型，那么可以认为模型已经将上下文关联考虑进去了，那么直接使用Greedy Search也是情有可原。

Beam Search

另外一种是beam search。常见的beam search思路是这样的

  1. 模型读入一帧MFCC，然后给出当下各种字符的概率
  2. Beam search利用这一层的概率展开搜索，然后取解空间中最优的k条路径/前缀，然后把这些前缀挨个输入到模型中

重复上面两个步骤就可以获得一个比较优秀的解。用伪代码表示就是

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def beam_search(y, beam_size):
  """
  param y:  模型的输出，形状为seq_len * n_class
  param beam_size:  搜索的宽度
  return:   最优的beam_size条路径
  """
  seq_len, n_class = y.shape
  logY = np.log(y)

  beam = [([], 0)]
  
  for t in range(seq_len):
    buff = []
    for prefix, score in beam:
      for i in range(n_class):
        new_prefix = prefix + [i]
        new_score = score + logY[t][i]
        buff.append((new_prefix, new_score))
    
    # 选取最优的beam_size个路径
    buff.sort(key=lambda x: x[1], reverse=True)
    beam = buff[: beam_size]
  
  return beam

但是当我用先前了解到的beam search知识去套参考代码（见参考）中的beam search时，发现事情并没有这么简单。

首先，参考代码里的beam search是直接输入模型在全部时间上的输出，也就是一个seq_len × n_class的矩阵，后续的搜索操作不再有模型的参与！这就导致一个问题，原本beam search中两个步骤交替保证了 后一个输入与之前所有的输入相关，但是这么一“偷懒”相关性给整没了。

用一个简单的例子说明一下，当模型预测第4个位置的输出时，如果`[h,e,l]`和`[r,e,l]`都在beam search的备选菜单中，那么模型下一步选择`[l]`的概率是不同的，因为`hello`很常见，但是`rell-`这样的组合很少见，所以`[h,e,l,l]`更能被模型认可。

如果取消了这个关联性，那么模型虽然可以选择k个较优解，但是其中概率最大的解一定和greedy search给出的解相同。带着这样的猜测我去运行了一下两种搜索算法看看它们的输出，结果让我大跌眼镜，两者的输出千差万别，可见这个版本的beam search暗藏玄机。

第二个不同点是，参考代码的beam search里还有两个不常见的变量，n_p_b（以blank结束的概率）和n_p_nb（不以blank结束的概率），并且这两个变量还存在前后状态之间的状态转移方程，也就是beam search中嵌套了一个动态规划算法。这么做的目的是什么呢？

Prefix Beam Search

这个将动态规划和常规beam search融合的算法称为prefix beam search。

首先说一下为什么有必要修改原来的beam search。假设字符集为小写字母加上一个空串ε，那么假设有三个串[h,h,ε,e,l]、[h,ε,e,e,l]和[h,ε,e,ε,l]，虽然它们所走的路径不同并各自拥有不同的概率，但是在CTC看来它们应当被视作同一种情况——[h,e,l]，因此这三者的概率需要相加。当然还有更多可能的路径都会被映射到前缀[h,e,l]，这些路径可能各自的概率不大，但是合在一起的概率就足以让模型确信该输出就是正确的结果。

我的版本

根据这个思路，我们很容易想到能否在beam search的基础上简单修改？我在了解了prefix beam search的动机后急于验证该算法的有效性，便很快实现了一个版本：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def prefix_beam_search(y, beam_size, blank):
    seq_len, n_class = y.shape
    logY = np.log(y)
    beam = [([], 0)]

    for t in range(seq_len):
        buff = []
        for prefix, p in beam:
            for i in range(n_class):
                new_prefix = list(prefix) + [i]
                new_p = p + logY[t][i]
                buff.append((new_prefix, new_p))

        new_beam = defaultdict(lambda: ninf)
        for prefix, p in buff:
            # norm_prefix的作用是将前缀化简，比如[1,1,_,2] ==> [1,2]
            # 但是结尾的blank会得到保留，防止出现歧义，比如[1,1,_] ==> [1,_]
            prefix = norm_prefix(prefix, blank)
            new_beam[prefix] = logsumexp(new_beam[prefix], p)

        # 选取最优的beam_size个路径
        new_beam = sorted(new_beam.items(), key=lambda x: x[1], reverse=True)
        beam = new_beam[: beam_size]

    return beam

然而当我将这个版本的算法和beam search做比较时，却发现两者的输出比较相似，并没有发挥出prefix beam search应有的作用。而官方版本的prefix beam search却能将输出压缩到合理的长度。

1
2
3
beam seach:   wihejvivleawewnin zlsqp tkjew
mine prefix:  agetvivleawewnin zlsqp tkjew
true prefix:  afkvqexhctlpdba

明明我使用norm_prefix将串压缩到了最简化，并留心到结尾的blank需区别对待了鸭？？难道官方的版本还有其他技巧？在此有必要将norm_prefix的实现贴出，并解释一下结尾的blank为何重要。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def norm_prefix(prefix, blank) -> tuple:
    """归一化前缀
    1. 去除开头blank，但保留结尾blank
    2. 去除结尾重复字符
    """
    if prefix[0] == blank:
        prefix.remove(blank)
    if len(prefix) >= 2:
        if prefix[-2] == prefix[-1]:
            prefix.pop()
        elif prefix[-2] == blank:
            prefix.remove(blank)

    return tuple(prefix)

如果当前的前缀是[a,b,_]，倘若去掉blank，那么当下一个字符是b时，将得到[a,b,b]，化简一下就变成了[a,b]，也就是失去了保留重复字符的能力。

这个问题困扰了我半天，不过幸好并最后发现了这个问题的症结所在：正是防止产生歧义，[a,b]和[a,b,_]被分开计算概率这个操作导致我自己的版本压缩能力差。而官方的版本[a,b]和[a,b,_]会被放在同一个beam node中。因此相比于官方版本我的版本在剪枝时丢掉了更多的路径。

官方实现

有了上面的发现，导出官方版本的prefix beam search也不是什么难事了，不如偷个懒直接给出代码，更详细的推导过程可以参考Sequence Modeling
With CTC。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def prefix_beam_search(y, beam_size, blank):
    T, V = y.shape
    log_y = np.log(y)
    beam = [(tuple(), (0, ninf))]

    for t in range(T):
        new_beam = defaultdict(lambda: (ninf, ninf))
        for prefix, (p_b, p_nb) in beam:
            for i in range(V):
                p = log_y[t, i]
                if i == blank:
                    new_p_b, new_p_nb = new_beam[prefix]
                    new_p_b = logsumexp(new_p_b, p_b + p, p_nb + p)
                    new_beam[prefix] = (new_p_b, new_p_nb)
                    continue
                end_t = prefix[-1] if prefix else None
                # 判断之前beam项中的最后一个元素和i的元素是不是一样
                new_prefix = prefix + (i,)
                new_p_b, new_p_nb = new_beam[new_prefix]
                # 如果不一样，则将i这项加入路径中
                if i != end_t:
                    new_p_nb = logsumexp(new_p_nb, p_b + p, p_nb + p)
                else:
                    new_p_nb = logsumexp(new_p_nb, p_b + p)
                new_beam[new_prefix] = (new_p_b, new_p_nb)
                # 如果一样，保留现有的路径，但是概率上要加上新的这个i项的概率
                if i == end_t:
                    new_p_b, new_p_nb = new_beam[prefix]
                    new_p_nb = logsumexp(new_p_nb, p_nb + p)
                    new_beam[prefix] = (new_p_b, new_p_nb)

        # 选取最优的beam_size个路径
        beam = sorted(new_beam.items(), key=lambda x: logsumexp(*x[1]), reverse=True)
        beam = beam[:beam_size]
    return beam

参考