58同城AI Lab在WeNet中开源Efficient Conformer模型

原创
2023/01/12 09:50
阅读数 428

1.模型介绍

2.模型实现

3.流式推理

4.实验结果


2022年8月,58同城TEG-AI Lab语音技术团队完成了WeNet端到端语音识别的大规模落地,替换了此前基于Kaldi的系统,并针对业务需求对识别效果和推理速度展开优化,取得了优异的效果,当前录音文件识别引擎处理语音时长达1000万小时/年,流式语音识别引擎支持语音对话量超过5000万次/年,详细工作可以参考《58同城:WeNet端到端语音识别大规模落地方案[1]》。

在优化工作中,我们复现了Efficient Conformer[2]模型,在实际场景数据上,与Kaldi最优模型相比,识别效果上CER绝对降低3%,解码性能提升61%。与Conformer相比,识别效果上CER从10.01%降低至9.30%,解码性能提升10%,结合int8量化,解码性能可提升60%。我们也在AISHELL-1公开数据集上进行了评测,CER为4.56%(No LM)。模型代码已开源至WeNet[3]

本文主要介绍我们对Efficient Conformer的复现工作,包含:模型介绍、模型实现、流式推理支持以及相关实验结果。


01

模型介绍

Conformer结构在语音识别领域取得了非常好的效果,已经被广泛应用于各种模型架构中。为了降低Conformer的计算复杂度、加快推理速度、减少所需的计算资源,Efficient Conformer对Conformer做出改进,提出了几种高效模型结构。实验结果证明,Efficient Conformer相比Conformer模型取得了更好的识别效果,以及更快的训练和解码速度。

Efficient Conformer的主要改进点如下:

  • Progressive Downsampling:Efficient Conformer Block在卷积模块中增加了下采样操作,降低时间维度,从而减小下采样后的(Efficient) Conformer Block的计算复杂度;

  • Grouped Attention:Efficient Conformer Block改进了Multi-Head Self Attention,增加grouped操作将自注意力模块的计算复杂度从O(n2d)降低为O(n2d/g),n为时间维度,d为隐层维度,g为group_size。

 


1.1 Progressive Downsampling

Efficient Conformer Encoder不同于典型Conformer,Conformer Block之前的下采样层使用1/2 subsampling(conv2d2)替换原始Conformer的1/4 subsampling(conv2d),其后一共包含三个stage,如上图右侧所示。前两个stage在N个Conformer Block之后叠加Downsampling Block,沿着时间维度进行下采样;最后一个stage叠加N个Conformer Block,不再叠加Downsampling Block。

Downsampling Block结构如下图所示,实现方式为将Conformer Block中的DepthwiseConv对应的stride设置为大于1的值,从而实现时间维度下采样。因为下采样后输出的shape比输入的shape小,因此残差模块需要增加Pointwise Projection模块将输入和输出映射到相同的维度。


1.2 Grouped Multi-head Self-Attention(Grouped MHSA)

传统Multi-head Self-Attention模块中Q、K、V大小为(n, d),该模块的计算复杂度为O(n2d);Grouped MHSA 首先将Q、K、V的维度变换为(n/g, d*g),其中g为group_size,再进行attention计算,最后将维度变换为原始的(n, d)。变换后,注意力模块的计算复杂度可降低为O(n2d/g)。


此外,作者还提出了Stride Multi-Head Self-Attention、

Relative Multi-Head Self-Attention 和 Local Multi-Head Self-Attention 等高效MHSA结构,感兴趣的朋友可以查阅原论文和代码。


02

模型实现

我们参考Efficient Conformer的代码[4]在WeNet开源项目上进行了复现,即在wenet文件夹下增加了efficient_conformer模块。


2022年12月26日我们向WeNet开源项目提交PR,贡献代码1426行,2023年1月4日被WeNet正式合并,主要开发者:周维、王亚如、李咏泽。

主要模块的实现细节如下:

2.1 Strided Convolution

(1)在depthwise_conv模块初始化定义时传入下采样步长stride参数,实现沿时间维度的下采样(convolution.py)。

self.depthwise_conv = nn.Conv1d(    channels,    channels,    kernel_size,    stride=stride,  # for depthwise_conv in StrideConv    padding=padding,    groups=channels,    bias=bias,)

(2)mask同步下采样(convolution.py):卷积模块增加下采样后,用于返回的输出数据的shape减小,对应返回的mask也应进行相同stride的下采样,以保持输出数据和mask的匹配。

# mask batch paddingif mask_pad.size(2) > 0:  # time > 0    if mask_pad.size(2) != x.size(2):        mask_pad = mask_pad[:, :, ::self.stride]    x.masked_fill_(~mask_pad, 0.0)

(3)带pointwise projection layer的残差结构(encoder_layer.py):由于卷积模块增加了下采样,导致输出维度小于输入维度,因此卷积模块对应的残差模块需要增加下采样操作。

# add pointwise_conv for efficient conformerif self.pointwise_conv_layer is not None:    residual = residual.transpose(1, 2)    residual = self.pointwise_conv_layer(residual)    residual = residual.transpose(1, 2)    assert residual.size(0) == x.size(0)    assert residual.size(1) == x.size(1)    assert residual.size(2) == x.size(2)

2.2 Grouped Multi-Head Self-Attention

(1)初始化中传入group_size参数,并重新定义位置编码偏置矩阵大小(attention.py)。

class GroupedRelPositionMultiHeadedAttention(MultiHeadedAttention):    def __init__(self, n_head, n_feat, dropout_rate, group_size=3):        """Construct an RelPositionMultiHeadedAttention object."""        super().__init__(n_head, n_feat, dropout_rate)        # linear transformation for positional encoding        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)        self.group_size = group_size        self.d_k = n_feat // n_head  # for GroupedAttention        self.n_feat = n_feat        # these two learnable bias are used in matrix c and matrix d        # as described in https://arxiv.org/abs/1901.02860 Section 3.3        self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k * self.group_size))        self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k * self.group_size))        torch.nn.init.xavier_uniform_(self.pos_bias_u)        torch.nn.init.xavier_uniform_(self.pos_bias_v)

(2)增加padding和reshape函数pad4group(attention.py):对Q、K、V、P在时间维度上根据group_size进行padding,保证可被group_size整除;padding会进行补0操作,不丢弃原始数据;padding之后即可按照group_size对Q、K、V、P进行维度变换;由于维度变换之后时间维度降低,同步地mask也需要降维,此处直接对mask下采样即可。

def pad4group(self, Q, K, V, P, mask, group_size: int = 3):    # Compute Overflows    overflow_Q = Q.size(2) % group_size    overflow_KV = K.size(2) % group_size
padding_Q = (group_size - overflow_Q) * int( overflow_Q // (overflow_Q + 0.00000000000000001)) padding_KV = (group_size - overflow_KV) * int( overflow_KV // (overflow_KV + 0.00000000000000001))
batch_size, _, seq_len_KV, _ = K.size()
# Input Padding (B, T, D) -> (B, T + P, D) Q = F.pad(Q, (0, 0, 0, padding_Q), value=0.0) K = F.pad(K, (0, 0, 0, padding_KV), value=0.0) V = F.pad(V, (0, 0, 0, padding_KV), value=0.0)
if mask is not None and mask.size(2) > 0 : # time2 > 0: mask = mask[:, ::group_size, ::group_size]
Q = Q.transpose(1, 2).contiguous().view( batch_size, -1, self.h, self.d_k * group_size).transpose(1, 2) K = K.transpose(1, 2).contiguous().view( batch_size, -1, self.h, self.d_k * group_size).transpose(1, 2) V = V.transpose(1, 2).contiguous().view( batch_size, -1, self.h, self.d_k * group_size).transpose(1, 2)
# process pos_emb P_batch_size = P.size(0) overflow_P = P.size(1) % group_size padding_P = group_size - overflow_P if overflow_P else 0 P = F.pad(P, (0, 0, 0, padding_P), value=0.0) P = P.view(P_batch_size, -1, self.h, self.d_k * group_size).transpose(1, 2)
return Q, K, V, P, mask, padding_Q

(3)forward_attention函数中(attention.py),在attention计算完毕后,输出数据需要变换为padding之后的维度,再去掉padding部分,截取有效输出。

# n_feat!=h*d_k may be happened in GroupAttentionx = (x.transpose(1, 2).contiguous().view(n_batch, -1, self.n_feat)     )  # (batch, time1, d_model)if padding_q is not None:    # for GroupedAttention in efficent conformer    x = x[:, :x.size(1) - padding_q]

2.3 其它细节实现

(1)pointwise projection(encoder.py

由于卷积模块进行了下采样,导致输入和输出维度不匹配,残差模块需要增加pointwise projection layer或下采样层将输入和输出维度进行统一。该模块在实现上有多种选择,比如可以使用卷积下采样或Pooling下采样。

我们在实验中对比了卷积下采样(kernel=3, stride=2, causal=true)和Pooling下采样(AvgPool1d)两种实现方式。实验结果显示,卷积下采样效果不如Pooling;且考虑到卷积的流式实现需要设置causal为true、并开辟和维护cache,而Pooling没有参数、不需要cache,简单适配即可直接用于流式训练和解码,最终我们采用AvgPool1d实现卷积模块的残差下采样。

# conformer module definitionif i in self.stride_layer_idx:    # conformer block with downsampling    convolution_layer_args_stride = (        output_size, self.cnn_module_kernels[index], activation,        cnn_module_norm, causal, True, self.stride[index])    layers.append(StrideConformerEncoderLayer(        output_size,        encoder_selfattn_layer(*encoder_selfattn_layer_args),        positionwise_layer(*positionwise_layer_args),        positionwise_layer(            *positionwise_layer_args) if macaron_style else None,        convolution_layer(            *convolution_layer_args_stride) if use_cnn_module else None,        torch.nn.AvgPool1d(            kernel_size=self.stride[index], stride=self.stride[index],            padding=0, ceil_mode=True,            count_include_pad=False),   # pointwise_conv_layer        dropout_rate,        normalize_before,        concat_after,    ))

(2)不同encoder layer attention维度不变(encoder.py

原始论文中,为了平衡下采样前后不同层的计算量,不同stage的Attention维度不同,且逐渐增大;我们的实现中为了简单方便灵活、节约计算资源、减小计算量,所有layer的Attention维度均保持一致;仅使用Grouped MHSA对不同layer的计算复杂度进行平衡,可以灵活设计添加Grouped MHSA模块的位置及group_size。通常选择连续对Downsampling Block之前的N层增加grouped操作。


03

流式推理

WeNet框架在流式推理时调用encoder中的forward_chunk接口,音频被划分为chunk_size大小输入到模型,在Conformer架构下实现流式需要记录Attention的K和V作为下一个chunk模型推理的att_cache。另外,流式场景下depthwize_conv通常使用因果卷积(Casual Convolution),因此需要记录隐变量的后 kernel_size-1 维向量作为cnn_cache,以便在下一个chunk卷积计算时使用,如下图所示。

而Efficient Conformer在Strided Convolution层会对输入向量在时间维度做下采样操作,导致cache在时间维度缩短,为了保证接口不变,需要在forward_chunk中对每层的cache做适当的“填充”。

我们首先参考Squeezeformer的实现方式,用以下函数计算每层的下采样倍数(encoder.py)。

def calculate_downsampling_factor(self, i: int) -> int:    factor = 1    for idx, stride_idx in enumerate(self.stride_layer_idx):        if i > stride_idx:            factor *= self.stride[idx]    return factor

对于att_cache,我们在时间维度重复factor倍数,并在使用att_cache时按照对应的factor进行下采样(encoder.py)。

for i, layer in enumerate(self.encoders):    factor = self.calculate_downsampling_factor(i)    xs, _, new_att_cache, new_cnn_cache = layer(        xs, att_mask, pos_emb,        mask_pad=mask_pad,        att_cache=att_cache[i:i + 1, :, ::factor, :],        cnn_cache=cnn_cache[i, :, :, :]        if cnn_cache.size(0) > 0 else cnn_cache    )    # shape(new_att_cache) = [batch, head, time2, outdim]    new_att_cache = new_att_cache[:, :, next_cache_start // factor:, :]    # use repeat_interleave to new_att_cache    new_att_cache = new_att_cache.repeat_interleave(repeats=factor, dim=2)

对于cnn_cache,则直接padding到 kernel_size-1 即可(encoder.py)。

# shape(new_cnn_cache) = [1, batch, outdim, cache_t2]new_cnn_cache = new_cnn_cache.unsqueeze(0)# padding new_cnn_cache to cnn.lorder for casual convolutionnew_cnn_cache = F.pad(    new_cnn_cache,    (self.cnn_module_kernel - 1 - new_cnn_cache.size(3), 0))

由于Grouped MHSA实际为reshape操作,所以new_att_cache的记录需要在reshape操作之前进行,这样可以保证cache的维度与常规MHSA一致(attention.py)。

if cache.size(0) > 0:    # use attention cache    key_cache, value_cache = torch.split(        cache, cache.size(-1) // 2, dim=-1)    k = torch.cat([key_cache, k], dim=2)    v = torch.cat([value_cache, v], dim=2)new_cache = torch.cat((k, v), dim=-1)
# May be k and p does not match. eg. time2=18+18/2=27 > mask=36/2=18if mask is not None and mask.size(2) > 0: time2 = mask.size(2) k = k[:, :, -time2:, :] v = v[:, :, -time2:, :]
# q k v p: (batch, head, time1, d_k)q, k, v, p, mask, padding_q = self.pad4group(q, k, v, p, mask, self.group_size)

forward_chunk接口会输入offset用于计算流式模式下的相对位置编码,由于Efficient Conformer会对时间维度下采样,导致输出与输入维度不匹配,因此通过y.size(1)计算的offset为下采样后的值,需要在使用时按照模型整体下采样倍数恢复(encoder.py)。

# using downsampling factor to recover offsetoffset *= self.calculate_downsampling_factor(self.num_blocks + 1)


04

实验结果

使用Efficient Conformer,只需要在配置文件中配置encoder参数即可。

encoder: efficientConformer

在自有场景下测试,Efficient Conformer获得了好于Conformer的效果,具体可参考文章[1]


同时我们在AISHELL-1上进行了两版模型实验。V1为我们线上使用的结构,即在Encoder的前1/3处使用Strided Convolution下采样(共12层),前4层均为Grouped MHSA结构,同时cnn_module_kernel会在Strided Convolution之后缩减同样倍数(如15->7),前置下采样使用1/4 subsampling的conv2d。V1 large将output维度从256增加至512,同时cnn_module_kernel从15增加至31。

    efficient_conf:        stride_layer_idx: [3]           # layer id with StrideConv        stride: [2]                     # stride size of each StrideConv        group_layer_idx: [0, 1, 2, 3]   # layer id with GroupedAttention        group_size: 3                   # group size of every GroupedAttention layer        stride_kernel: true             # true: recompute cnn kernels with stride

V2符合原始Efficient Conformer论文结构,前置下采样为1/2 subsampling的conv2d2,在Encoder的1/3和2/3处做两次下采样,cnn_module_kernel固定不变。

    efficient_conf:        stride_layer_idx: [3, 7]           # layer id with StrideConv        stride: [2, 2]                     # stride size of each StrideConv        group_layer_idx: [3, 7]            # layer id with GroupedAttention        group_size: 3                      # group size of every GroupedAttention layer        stride_kernel: false               # true: recompute cnn kernels with stride

在不使用语言模型的情况下效果如下:


可见,在WeNet项目当前已合并的模型中(截止2023-1-9),Efficient Conformer在AISHELL-1公开数据集上目前最优的CER为4.56%,超过Conformer效果4.61%。

后续计划:

(1)进一步完善开源数据上的测试效果

(2)支持Efficient Conformer的ONNX导出,和GPU流式部署


参考文献

[1] 58同城:WeNet端到端语音识别大规模落地方案

[2] Efficient Conformer: https://arxiv.org/pdf/2109.01163.pdf

[3] WeNet Efficient Conformer PR:https://github.com/wenet-e2e/wenet/pull/1636

[4] Efficient Conformer Code: https://github.com/burchim/EfficientConformer


作者介绍:

周维,58同城TEG-AI Lab算法架构师,语音算法部负责人,负责语音识别、语音合成算法研发。

王亚如,58同城TEG-AI Lab语音算法部算法高级工程师,主要负责端到端语音识别算法研发。


58同城AI Lab部门简介

58同城AI Lab隶属TEG技术工程平台群,旨在推动AI技术在58同城的落地,打造AI中台能力,以提高前台业务人效、收入和用户体验。


推荐阅读

3人半年打造语音识别引擎——58同城语音识别自研之路

PPT+视频回放 | 语音技术在58同城的应用

人物 | 王焱:58同城流式语音识别引擎应用实践

流式和离线语音场景下VAD语音端点检测算法实践

沙龙干货 | 58同城语音识别技术的探索和实践

语音识别中的WFST和语言模型

语音识别中Chain Model的原理和实践

本文分享自微信公众号 - 58技术(architects_58)。
如有侵权,请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”,欢迎正在阅读的你也加入,一起分享。

展开阅读全文
加载中
点击引领话题📣 发布并加入讨论🔥
打赏
0 评论
0 收藏
0
分享
返回顶部
顶部