《JETS Jointly Training FastSpeech2 and HiFi-GAN for End to End Text to Speech》_怎么训练jets模型

《JETS: Jointly Training FastSpeech2 and HiFi-GAN for End to End Text to Speech》

1. 背景

现在比较成熟的TTS系统一般包括两个模型：声学模型和声码器。前者旨在根据输入文本生成声学特征，例如filter bank；后者旨在从声学特征恢复语音波形。此时可能存在声学特征mismatch的问题，从而影响合成语音质量。具体来说，声码器在训练时使用golden声学特征作为输入，而在推理时使用声学模型预测的声学特征作为输入。

为了解决声学特征mismatch的问题，有以下方案可供参考。

1. 1）2）同时开始训练，再用 3）来缓解mismatch问题

   1）训练声学模型

   2）使用golden声学特征训练声码器

   3）以 2）作为预训练模型，使用1）预测的声学特征做finetune

2. 先训练声学模型，再基于预测的声学特征训练声码器。

虽然以上方案可以解决mismatch问题，但训练流程比较复杂。方案1需要额外的finetune步骤，方案2的声码器训练只能等声学模型训练结束才可以开始。

本文所提方案简化了上述过程，并且通过加入对齐学习避免了使用额外的对齐模型。

PS：在本文之前已经出现过类似想法的方案，这里暂不涉及。

2. 方案

本文所提方案主要基于FastSpeech2和HiFi-GAN模型，此外还有一个对齐模块。这里不详细介绍FastSpeech2和HiFi-GAN，着重介绍如何实现联合训练。

如图1所示，Generator之前是FastSpeech2模块，即声学模型；之后则是HiFi-GAN模块，即声码器（这里省略了HiFi-GAN的discriminator模块）。根据前文提到的内容，Decoder的输出是声学特征，文章说在这里不会计算声学特征损失，从而实现文本到波形的tts系统。

图中还有一个Alignment模块，它的输入是文本和mel谱。作用是实现文本音素个数和mel谱帧数的函数关系，为模型训练提供每个音素的时长信息。一般一个文本音素会对应多个mel帧，文中提到会使用CTC损失实现两者的对齐。细节上还没搞懂，对比代码应该会更容易理解 。（该模块只在训练时使用，目的是给模型提供音素的golden时长信息）

前面已经说了Alignment模块的作用是为音素提供golden时长。此过程中涉及到3个公式，如下：

其中，h 是text embeddings过2个卷积层得到的向量，m是mel谱过3个卷积层得到的向量。i 和 j 分别代表音素个数和mel帧个数。D(i,j)表示不同 h 和 m 之间的 L2距离，然后在 h 维度上做softmax 计算出当前帧与不同 h 向量的softmax概率值。

熟悉这个公式的应该会发现，该公式就是CTC loss。等式右边的乘积表示一种不同帧音素和帧对应关系的组合，加和则表示所有的这种可能性（用CTC loss的话则是乘积表示一种可能的路径，加和表示所有路径的集合）。因为CTC loss要求输入比输出要长，即需要满足多对一的映射关系，因此该公式可以认为是求一种极大似然，值越大表示可能性越高，是正确关系的概率越大。作为损失函数，需要对其取负值，即 loss_forwardSum = -P(S(h)|m)
还有一点是此公式中的S表示什么，S表示文本和谱有效对应关系的集合。CTC loss中有提到，由于输入输出满足多对一的关系，因此存在太多可能性，导致其复杂度过高。因此会采用一个动态规划的方式只将有效的对应关系（来自A_soft值）加入到S集合中。具体一种对应关系是什么形式呢。举个例子

假设有 3个音素，共对应5帧mel谱
H：h1, h2, h3, h3, h3
M：m1, m2, m3, m4, m5
S：s1: h1=m1, s2: h2=m2, s3: h3=m3, s4: h3=m4, s5: h3=m5
即第1，2个音素分别对应一个mel帧，第3个音素对应3个mel帧。S表示h和m的这种对应关系。
根据公式，假设t=3,即
P(s3|m3)表示当条件是m3时，关系是s3: h3=m3，即第3个音素对应第3帧的概率
1
2
3
4
5
6
7

当找到最可能的对应关系后，可根据S拿到最可能的路径。但是如何拿到时长信息呢，论文使用monotonic alignment search (MAS)将软对齐分布A_soft转为单调二值化的硬对齐分布A_hard，该分布满足两个条件：1）token个数与音素个数相同，2）所有token的时长相加等于mel帧数。通过计算两个分布之间的KL散度使得A_hard分布接近A_soft分布。

Variance adaptor即变量适配器，主要用来预测时长、基频和能量信息。其中，时长信息是每个音素持续的帧数；基频和能量对于提高合成音频质量很重要。因此这里会分别计算时长、基频和能量损失。图2中的LR（length regulator），在文中使用高斯上采样实现文本音素在时间维度上的expand。例如 text embeddings的shape是 [12, 128] ，表示该文本有12个音素。给定每个音素的时长信息，则可以对每个音素进行扩展，将其变为 [mel_len, 128] 形状的向量，此时所有音素持续的时间相加就是mel_len。其实此处也可以使用重复上采样，但是这种方式合成的音频会有点生硬。高斯上采样会相对流畅些，合成音频有更好的听感。

高斯上采样的具体细节可以参考《Non Attentive Tacotron》一文。

综上所述，最终的损失函数有：loss_all = loss_gan + w1 * loss_variance + w2 * loss_align
其中HiFi-GAN的对抗损失：loss_gan = loss_generator + loss_dis + w_fm * loss_fm + w_mel * loss_mel
时长、基频、能量损失分别计算均方误差。

实验相关内容请自行阅读论文。值得一提的是，文中设置的对比实验完全覆盖了背景中提到的两种解决mismatch问题的方案。
想吐槽一下，论文没有给出详细的训练过程。如果不引入预训练模型，感觉需要训练很长时间。而且对于Alignment模块的实现细节如果不对照代码学习，很难琢磨透。