MB-iSTFT-VITS 模型论文思路与实验分享：基于VITS架构优化的轻量级文本转语音模型

参考文献：

[1] Kawamura M, Shirahata Y, Yamamoto R, et al. Lightweight and high-fidelity end-to-end text-to-speech with multi-band generation and inverse short-time fourier transform[C]//ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2023: 1-5.

[2] Ren Y, Ruan Y, Tan X, et al. Fastspeech: Fast, robust and controllable text to speech[J]. Advances in neural information processing systems, 2019, 32.

[3] Kim J, Kong J, Son J. Conditional variational autoencoder with adversarial learning for end-to-end text-to-speech[C]//International Conference on Machine Learning. PMLR, 2021: 5530-5540.

[4] Kong J, Kim J, Bae J. Hifi-gan: Generative adversarial networks for efficient and high fidelity speech synthesis[J]. Advances in neural information processing systems, 2020, 33: 17022-17033.

[5] Kaneko T, Tanaka K, Kameoka H, et al. iSTFTNet: Fast and lightweight mel-spectrogram vocoder incorporating inverse short-time Fourier transform[C]//ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2022: 6207-6211.

本文内容主要来源于此论文：Demonstration of Lightweight and High-Fidelity End-to-End Text-to-Speech with Multi-Band Generation and Inverse Short-Time Fourier Transform - Masaya Kawamura

该论文主要针对VITS模型做了轻量化架构优化，使用了逆短时傅里叶变换（iSTFT）和多频段生成技术，替代了原VITS模型中解码器的部分上采样层，在仅损失较小的语音合成质量的同时，大大加速了语音生成的速度，加速比大约为3.4倍 [1]。本博客将介绍该模型的主要工作、模型架构优化思路，以及论文作者和笔者做的部分相关实验结果，并在原论文的基础上添加部分背景知识介绍。

一、基础模型的选择

我们一般会根据生成语音序列的方式，将端到端文本转语音（TTS）模型分为自回归模型和非自回归模型。

• 自回归模型是逐步生成语音序列的模型。它们根据前一步生成的语音帧或特征来预测下一步的输出，我们熟知的 Tacotron、Deep-Voice、Transformer-TTS 等都是自回归模型。
• 非自回归模型是并行生成语音序列的模型。它们不依赖于前一步的输出，而是可以同时生成多个步骤的输出，FastSpeech 系列、NaturalSpeech 系列以及本文重点 VITS 等都是非自回归的。

很明显，非自回归能够并行生成整个语音序列的特性，使得其推理速度要远远快于自回归模型，速度大约差2~3个数量级。下面给出了经典的非自回归模型 FastSpeech 和自回归模型 Transformer TTS 的语音合成速度对比图 [2]：

鉴于文章的 Motivation 是要做一个轻量级文本转语音模型，推理速度当然是重中之重。因此基本敲定了选用非自回归模型作为基础模型。而又因为 VITS 模型合成的语音无论是在自然度还是清晰度上都非常好，语音质量非常高，因此文章选择了 VITS（一代） 作为基础模型。实际上，很多最近一段时间大火的 TTS 模型，比如 Bert-VITS2、GPT-SoVITS 等，都是基于 VITS。

笔者接下来也将开始对GPT-SoVITS做相关学习与研究，敬请期待！

二、模型架构优化思路

如果有读者想深入了解 VITS 模型架构，也可以参考笔者的文章：VITS 模型详解与公式推导：基于条件变分自编码器和对抗学习的端到端语音合成模型_vits tts 模型展开-CSDN博客。这里将略述 VITS 的架构。

VITS 推理瓶颈分析

既然我们需要优化 VITS，我们就需要知道 VITS 在推理时其瓶颈究竟在哪里。我们简单看一下 VITS 的推理架构图 [3]：

根据架构图，我们可以把 VITS 分为以下几个部分：

• 文本编码器（Text Encoder）
• 随机持续时长预测器（Stochastic Duration Predictor）
• 注意力加权计算（Attention Computation）
• 逆向标准化流（Inverse Normalizing Flow）
• 解码器（Decoder）

在原论文中，作者直接给了实验结果。笔者则是使用官方提供的基于 LJ Speech 数据集的预训练模型，并随机在 Youtube 网站中抽取了100条长短不一的英文语句进行文本转语音任务，使用 python 的 time 库进行计时。具体在 infer 函数中的代码如下：

def infer(self, x, x_lengths, sid=None, noise_scale=1, length_scale=1, noise_scale_w=1., max_len=None):
    start_time = time.time()                                                                                                    # 计时
    x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths)
    text_enc_time = time.time()                                                                                                 # 计时
    # speaker embedding
    if sid is None: sid = torch.LongTensor([0]*x.shape[0])
    g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
    dp_time_s = time.time()
    # predict alignments
    logw = self.dp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) if self.use_sdp else self.dp(x, x_mask, g=g)
    dp_time_e = time.time()
    w = torch.exp(logw) * x_mask * length_scale
    w_ceil = torch.ceil(w)

    y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()  
    y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(x_mask.dtype)
    attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
    attn = commons.generate_path(w_ceil, attn_mask)

    m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']
    logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']

    z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
    attn_time_ = time.time()
    z = self.flow(z_p, y_mask, g=g, reverse=True)
    iflow_time_ = time.time() 
    o = self.dec((z * y_mask)[:,:,:max_len], g=g)
    end_time = time.time()
    elapsed_time = end_time - start_time
    text_enc_time = text_enc_time - start_time
    dp_time = dp_time_e - dp_time_s
    attn_time = attn_time_ - dp_time_e
    iflow_time = iflow_time_ - attn_time_
    decoder_time = end_time - iflow_time_
    return o, attn, y_mask, (z, z_p, m_p, logs_p), elapsed_time, text_enc_time, dp_time, attn_time, iflow_time, decoder_time

很简单，就是时间的加加减减而已。然后再写一个批量推理脚本，输出最后的推理占比平均值即可。最后笔者的实验结果如下：

而原论文的作者切分的部分比笔者的要少，主要只分为了文本编码器、标准化流、解码器三个部分，并且使用的也是实时率（Real Time Factor, RTF）作为实验结果，如下：

RTF 的概念我们之后会介绍，我们的结果和作者的结果在某种程度上高度相似：即解码器占用了整个推理过程的绝大部分时长。因此我们着手对 VITS 的解码器进行优化。

VITS 解码器架构优化

我们来回顾一下 VITS 的解码器部分。在原论文中，VITS 的解码器架构基于 HiFi-GAN V1[4] 的生成器（Generator），它主要由转置卷积（Transposed Convolution）堆栈组成，每个模块后面都有一个多感受野融合模块（multi-receptive field fusion, MRF）。 MRF 的输出是具有不同感受野大小的残差块的输出之和。就这样使用基于重复卷积的网络对输入声学特征进行上采样。因此我们首先考虑减少该模块中的冗余。

论文作者根据这个模块的主要目的是为了重建时域波形，并从一个声码器 iSTFTNet[5] 获得灵感，从而想到了逆短时傅里叶变换（Inverse Short-Time Fourier Transform，iSTFT）技术。声音在处理的时候通常会使用短时傅里叶变换将声音序列从时域转为频域，得到频谱（Spectrogram）和相位（Phase），那我们也可以使用逆变换，提供频谱和相位，从而合成语音。简单易懂的思路。

因此我们期望利用逆短时傅里叶变换去替代原来Decoder中的一些输出侧层，从而对Decoder架构进行简化，降低计算成本。笔者根据论文作者提供的代码，画出了简化示意图如下：

此外，由于这个过程酷似声码器对于时域波形的重建，而在许多声码器中都采用了多频段并行生成（multi-band parallel generation strategy）的策略，因此我们也将尝试使用多频段生成策略，与 iSTFT 进行组合，共同优化波形重建过程。

在声学中，频带是指在声音频谱中特定的频率范围。声音的频谱通常被分成多个频带，每个频带代表了一定范围内的频率，而人类的听觉系统一般对于不同频率范围内的声音有不同的感知特性。

而声码器中，多频带生成就是利用了神经网络的稀疏性，使用单个共享网络来生成所有子带信号，然后将这些频带的声音合并成最终的合成语音。这在保持合成质量的同时显着降低了计算成本。

通过组合 iSTFT 和多频带生成策略，我们最终可以给出改进后的Decoder架构，如下图所示。

• 首先，潜在变量 z 通过部分转置卷积和主要由残差块组成的多感受野融合模块进行上采样，通过一个输出卷积网络后，投影到每个子带信号的幅度和相位变量。
• 然后，运用 iSTFT，将幅度和相位变量组合，从频域重建时域波形，生成每个子带信号。
• 最后，我们会在样本之间添加零来对这些子带信号进行上采样，以匹配原始信号的采样率，然后使用固定的合成滤波器组将其集成到全带波形中。这里我们使用的合成滤波器是基于伪正交镜像滤波器组（pseudo-quadrature mirror filter bank, pseudo-QMF）的固定（不参与训练）合成滤波器。之后作者也给出了多频带结构中的可训练合成滤波器，用来提高语音合成质量。

自然的，由于添加了新的结构，我们的损失函数也要有所变动。除了以下损失函数外：

• 与目标波形的梅尔频谱重构损失（Reconstruction Loss）
• 潜在变量z的后验分布与先验分布的KL散度（KL Loss）
• 音素持续时间预测器训练损失（Duration Loss）
• 对抗学习的最小二乘损失函数（Adversarial Loss）
• 对抗学习额外的特征匹配损失函数（Feature-Matching Loss）

我们还额外添加了多分辨率短时傅里叶变换子带信号损失（Multi-resolution STFT Loss in sub-band signals），用于测量多频段子带信号和真实子带信号之间的差距。并且为了生成根据输入波形计算子带 STFT 损耗所必需的真实子带信号，我们同样使用基于伪正交镜像滤波器组的分析滤波器来进行生成。因此，基于上述的优化思路，最终的整个模型架构我们给出示意图如下：

论文作者为这个模型取名为MB-iSTFT-VITS。值得一提的是，为了稳定音素持续时间的预测，论文作者这里采用的是确定性持续时间预测器，而不是随机持续时间预测器。其实我感觉这里可能也是为了加快推理速度，不过我没有做过二者的推理时间差异实验，不太清楚这样的选择会对语音质量和推理速度造成什么样的影响。

三、模型测试实验

数据集选择与模型训练

由于原论文提供的代码不支持中文，因此我参考中文的 text_cleaner，对原论文做了中文适配，其实也就是写代码将中文文本转为音素（Phoneme）。然后原论文的代码也不支持多语者训练，不过好在他是基于原版 VITS 代码进行开发的，原版 VITS 支持多语者训练，所以他的代码也残留了些许苗头，我也做了相应的适配。

训练数据集上，由于笔者算力限制，只能选用很小的数据集进行少轮次训练。又由于最近喜欢搞情感，我就用了 Emotional Speech Database 的部分中文语者语音作为训练数据，从10个中文语者中，选择了声音音色具有一定明显特质的语者2和语者5。此外，有关训练集和测试集的划分，我们从每个语者的每个情感中的350条平行语句中抽取10条语句，共100条语句作为测试集。

训练数据集的具体情况如下：

我们把每个语者的每种情感都当成一个 Speaker 来进行训练。说实话，我也感觉训练语料实在是太少了，但是就像这样少的训练语料，只训练2000轮次（原论文中训练了20000个轮次），用一张3090卡也跑了109小时。这里给出重建损失（Reconstruction Loss）图，模型正常收敛了。

最后还用了 Gradio 简单写了个前端用来展示。我也是听别人介绍才知道了这个 AI 模型展示组件，确实超级好用，啥模态的模型我感觉都能用这个组件展示。最后合成的语音效果也只能说差强人意吧，毕竟训练语料和轮次摆在那里，也不期望有多么惊艳了，倒是得益于 VITS 强大的建模能力，也能听出来语音里面的情感。

模型测试实验具体介绍

首先我们来看看语音合成的质量如何。这里就直接放原论文中作者做的实验结果了，作者使用的是ONNX版本的：

这里面 MB-iSTFT-VITS 就是咱们的主角。MOS（平均意见分数）代表了语音合成质量，Params 表示参数数量，RTF 是实时率也就代表了合成速度。这里也包括了使用可训练滤波器的 MS-iSTFT-VITS、只使用 iSTFT 而没有使用多频段生成策略的 iSTFT-VITS ，以及它们的 mini 版本（单纯砍参数）的相关结果。感兴趣的读者可以看一下原文。

当然这里是英文语音的合成结果哈，不过我也使用了相同的训练语料、相同的中文 text_cleaner，在原版 VITS 的基础上也进行了训练，训练轮次为1000次（训练时长36小时），也用了我的中文 MB-iSTFT-VITS 的1000轮次模型文件进行语音合成对比，也给了很多同学听，基本上也是大差不差，仅比原版的 VITS 要低了一点，也简单算了一下 MOS：

然后就是重中之重，推理速度的实验对比了。原文作者通过 RTF 衡量模型推理速度，结果也可以参考上面的表格。RTF 简单介绍一下，其实就是一个模型生成一个单位时间的语音所需要的时长。其计算公式如下：
$$RTF=\frac{T_{syn}}{T_{wav}}$$
其中，T_syn为模型生成一段语音所需要的时间，通常以秒为计量单位；T_wav为当次合成的语音的长度（时长），同样也以秒为计量单位。RTF通常以1为临界点，当一个模型在一个测试集上的RTF为1时，说明它正好可以进行实时语音合成。RTF的值越低，说明模型合成语音速度越快，越轻量级。

笔者的实验则是在本地进行，本地环境如下：

用CPU推理，使用1000轮次的原版 VITS 和 MBVITS（笔者称呼），使用同样的训练数据集和中文文本预处理，能相同的参数都相同，从QQ群里面随机抽了110条中文文本进行语音合成，做了相关统计图如下：

其中，蓝色柱状图代表此条文本的长度，使用左边纵坐标轴进行衡量。红色的折线代表MBVITS的推理耗时，黄色的折现代表原版VITS的推理耗时，二者的单位均为秒，使用右边纵坐标轴进行衡量。从统计图中我们不难看出，MBVITS在所有语句上的推理耗时都低于原版VITS，这个比例（MBVITS耗时 : 原版VITS耗时）在0.1~0.6之间波动。

我也顺便算了 RTF，做了各模块的时长占比统计，统计表如下：

从统计图表中可以看出来，新模型的推理速度还是非常快的，在我这里的加速比为3.4倍，也和原论文中作者给出的结果相同。然后我们也补充一个消融实验，不使用多频段生成策略，也做类似的统计图：

其中，除了之前我们使用红色表示MBVITS，使用黄色表示原版VITS外，我们新添加了蓝色用于表示iSTFT-VITS。我们再将统计的三个模型的测试集推理时间、生成语音时长、RTF，并汇总成表格如下表

实际上从图表上来看，基于单 iSTFT 的轻量化效果确实已经非常好了，相较于原版 VITS 的加速比大约为1.9倍，但当我们将多频段生成策略和 iSTFT 相结合后，还会在单 iSTFT 的基础上再次加速大约1.8倍，因此这一个消融实验证明，多频段生成策略对模型优化确实起到了很大的帮助。

四、总结

我自己也在这一方面没啥研究，也不知道这个模型实验结果相对于其他模型是好是坏，反正我自己个人感觉还是很不错的。不过实际上我做的复现还远远不及原作者的工作，首先我还是在pytorch上面跑的代码，还是类似实验环境，原作者则是转成了ONNX的，我计划之后也试着把我的代码转成 NCNN，并在树莓派上部署一下看看啥情况，毕竟咱们做的是轻量化模型，肯定要在嵌入式设备上跑的，后续也打算写写博客介绍一下我第一次进行嵌入式设备部署的过程。然后就是这篇论文实际上只提高了推理速度，而没有减少模型大小，作者在实验中只是单纯直接砍了参数，这样倒也行，不过也有很多其他的轻量化手段我不知道可不可以用在这上面试着减少一下模型的大小。然后就是情感嵌入，我没做情感嵌入，这也是一大遗憾。最后，欢迎读者分享自己的理解与看法，本人学识尚浅定有疏漏，欢迎斧正。