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

参考