SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sungwon Kim $^{1,2*}$ , Kevin J Shih $^{1}$ , Rohan Badlani $^{1}$ , João Felipe Santos $^{1}$ , Evelina Bhakturina $^{1}$ , Mikyas Desta $^{1}$ , Rafael Valle $^{1}$ , Sungroh Yoon $^{2,3}$ , Bryan Catanzaro $^{1}$
NVIDIA<sup>1</sup>  
Department of Electrical and Computer Engineering, Seoul National University<sup>2</sup>  
Interdisciplinary Program in Artificial Intelligence, Seoul National University<sup>3</sup>
While recent large-scale neural codec language models have shown significant improvement in zero-shot TTS by training on thousands of hours of data, they suffer from drawbacks such as a lack of robustness, slow sampling speed similar to previous autoregressive TTS methods, and reliance on pre-trained neural codec representations. Our work proposes P-Flow, a fast and data-efficient zero-shot TTS model that uses speech prompts for speaker adaptation. P-Flow comprises a speech-prompted text encoder for speaker adaptation and a flow matching generative decoder for high-quality and fast speech synthesis. Our speech-prompted text encoder uses speech prompts and text input to generate speaker-conditional text representation. The flow matching generative decoder uses the speaker-conditional output to synthesize high-quality personalized speech significantly faster than in real-time. Unlike the neural codec language models, we specifically train P-Flow on LibriTTS dataset using a continuous mel-representation. Through our training method using continuous speech prompts, P-Flow matches the speaker similarity performance of the large-scale zero-shot TTS models with two orders of magnitude less training data and has more than  $20\times$  faster sampling speed. Our results show that P-Flow has better pronunciation and is preferred in human likeness and speaker similarity to its recent state-of-the-art counterparts, thus defining P-Flow as an attractive and desirable alternative. We provide audio samples on our demo page.
Zero-shot TTS (Text-to-Speech) refers to the goal of generating text-conditioned speech with novel speaker characteristics at inference without the need for additional training. The speaker characteristics are provided by a short reference audio segment, which in recent works can be as short as 3 seconds long [37, 19, 6]. Many recent advances have come at a cost: larger datasets, more complicated training setups, additional quantization steps, additional pretraining tasks, and computationally expensive autoregressive formulations. Our work explores whether a more efficient solution exists. By staying within the general framework of prompt continuation, as often seen in recent large language model works, we demonstrate that similar results can be achieved with simpler training pipelines, significantly less data, and faster inference.
Following recent successes in the zero-shot and few-shot capabilities of large scale language models such as GPT [2], it has been generally accepted that an analogous approach will yield similar results
* Work done as a research intern at NVIDIA.
† Corresponding authors: Sungwon Kim: ksw0306@snu.ac.kr, Rafael Valle: rafaelvalle@nvidia.com, Sungroh Yoon: sryoon@snu.ac.kr
in speech synthesis. However, it is not yet clear that recent approaches in large scale representation learning have a significant advantage over existing and cheaper traditional representations such as mel-spectrograms. Of note, recent works such as SPEAR-TTS and VALL-E make use of pretrained quantized neural audio codec models such as Encodec and SoundStream [10, 40]. The costs of using such a model are expensive. Feature extraction is significantly more computationally demanding than extracting mel-spectrograms from audio. Encodec in particular contains autoregressive components in both the encoding and decoding, which may be an issue where low-latency is necessary. Further, having too many different versions of such a codec inhibit interoperability across models. While it is likely that neural codecs may yet have significant advantages in the near future, our work demonstrates competitive results with standard mel-spectrograms, suggesting that the costs of neural codecs in the zero-shot TTS setting is not yet justified by any significant improvements in generalization.
Improving inference speed is the other focus of this work. Our work improves upon the latency of prior approaches by reducing reliance on autoregressive formulations, and by employing recent advances in ODE-based generative modeling such as Flow Matching [23, 24, 35]. In addition to autoregressive components within the neural codec architectures, prior works such as SPEAR-TTS and VALL-E also employ autoregressive decoders to infer audio token sequences. While autoregressive modeling is a natural choice for sequential data such as audio signals, our work demonstrates that a lower-latency non-autoregressive formulation is both viable and competitive. Another common computational bottleneck is the incorporation of probabilistic generative models such as denoising diffusion models in recent works. Such architectures exist in many recent state-of-the-art speech and audio synthesis models, but come at the cost of requiring numerous function evaluations during inference. Distillation techniques exist to reduce the number of function evaluations, but require additional training stages, further complicating the training pipeline. In contrast, we experiment with recently proposed flow-matching algorithms. These models are closely related to diffusion models, but actively encourage simpler and straighter trajectories, which in turn require significantly fewer function evaluations during inference – all without the need for additional distillation stages.
Our work proposes P-Flow, high-quality, data-efficient, and fast zero-shot TTS model with emergent in-context learning capabilities for speaker adaptation. Inspired by recent success using language-model prompting to achieve zero-shot TTS results, we introduce a speech-prompted text encoder within a non-autoregressive TTS model. P-Flow combines this prompted text encoder with a flow-matching generative decoder to efficiently sample high-quality mel-spectrograms. Our speech-prompted text encoder incorporates a 3-second random segment of speech as a prompt along with the text input and learns to predict the mel-spectrogram using the speech-prompted text input. Given the speaker-conditional output from the text encoder, our flow-matching decoder is trained to predict a vector field that models the probabilistic distribution of mel-spectrograms conditioned on the text encoder's output. Our prompting approach directly utilizes random segments of speech as input, which poses a risk of P-Flow learning trivial solutions such as the identity function for the random segment. We address the potential problem by masking the loss for the random segment during training. Our contributions are as follows:
- We propose a speech prompt approach for the non-autoregressive zero-shot TTS model which surpasses the speaker embedding approach and provides in-context learning capabilities for speaker adaptation.
- We propose a flow matching generative model for a high-quality and fast zero-shot TTS that significantly improves the synthesis speed and sample quality compared to the large-scale autoregressive baseline.
- We demonstrate comparable speaker adaptation performance to the large-scale autoregressive baseline using significantly fewer training data and a small transformer-based encoder highlighting the effectiveness of the proposed speech prompting approach.
- P-Flow achieves an average inference latency of 0.11 seconds on an NVIDIA A100 GPU.
Zero-Shot Speaker Adaptive TTS (Zero-Shot TTS): Instead of specifying speaker identity with discrete embedding vectors as in the case of multi-speaker TTS models, zero-shot TTS formulations assume the ability to extract speaker embedding equivalents on-the-fly from short audio samples. Common approaches either use a separate audio-sample-to-embedding encoder to extract speaker
vectors on the fly [17, 27, 20, 7, 3, 38], or formulate the problem as a continuation task with language model architectures and probabilistic generative models. Methods such as VALL-E and SPEAR-TTS [37, 19] employ the prompting paradigm from large transformer-based language models. Notably, they use the reference audio as the prompt, from which the model is expected to decode the most likely continuation. The decoding procedure is further conditioned with the provided text transcript to achieve a text-to-speech formulation. Relatedly, approaches such as [28, 9, 25] use the data inpainting paradigm from generative probabilistic models, treating the speech-to-be-synthesized as the missing data to continue the available data (prompted audio). Just as in the LM-based approach, the inpainting is guided by conditioning on the target text sequence. Our work operates in the same prompt-based zero-shot TTS setting as VALL-E and SPEAR-TTS, establishing competitive results with much less data, and additionally providing faster inference speeds via flow-matching-based generative models and non-autoregressive formulations.
Perhaps most architecturally related to our work is  $\mathrm{A}^3\mathrm{T}$  [4], which similarly trains a transformer ( $\mathrm{A}^3\mathrm{T}$  uses a Conformer [12]) model to infer text-conditional mel spectrograms with a masked reconstruction loss. However,  $\mathrm{A}^3\mathrm{T}$  uses a forced aligner, requires a text transcript for the zero-shot audio prompt, and focuses primarily on the pretraining task and audio reconstruction, with limited results in the zero-shot TTS domain. In contrast, our proposed approach uses an end-to-end audio-text alignment mechanism, uses short 3 second audio-only prompt, and focuses primarily on the zero-shot TTS task.
Generative Sampling with Fewer Steps: Generative models in the score matching and denoising diffusion model family typically require hundreds of iterative steps. It is common to reduce the number of steps required via approaches such as the DDIM [32] framework and distillation [31]. Notably, distillation will progressively map every two steps to a single step, halving the number of total steps required with every iteration of the algorithm. Recent works such as flow matching [24, 23, 35] and consistency training (from scratch) [33] offer a single-stage approach, encouraging approximately straight, causal trajectories in the initial model, thereby reducing the need for further distillation stages. Our work serves as a demonstration of the use of such approaches as flow matching in large-scale practical audio applications.
Our work aims to provide in-context learning capabilities for zero-shot speaker adaptation in a high-quality and fast non-autoregressive TTS model. To avoid the potential bottleneck of speaker embedding approaches to extract speaker information from the reference speech, we adopt a prompting approach that directly utilizes the reference speech as a prompt for speaker adaptation, similar to neural codec language models. In addition to a speech prompting method, we use a flow matching generative model as our decoder for efficient sampling. We provide an overview of our method in Section 3.1, followed by a detailed explanation of our decoder in Section 3.2, and we describe the remaining details in Section 3.3.
3.1 P-Flow
P-Flow training is similar to that of masked-autoencoders. Given <text, speech> paired data, we denote the mel-spectrogram of the speech as  $x$ , and the text as  $c$ . Let  $m^p$  be a indicator mask on the sequence  $x$ , masking out a randomly positioned 3 second segment from  $x$ . We define  $x^p = (1 - m^p) \cdot x$ . We use a  $p$  superscript here indicating that this variable will be replaced with arbitrary 3-second prompts during zero-shot inference. The training objective for P-Flow is as follows: reconstruct the  $x$  given  $c$  and some segment  $x^p$ . Note that even though we expect P-Flow to reconstruct the entirety of  $x$  during training, including the provided  $x^p$ , the model does the exact positioning of  $x^p$  within  $x$ .
P-Flow learns to model the conditional probability distribution of speech  $p(x|c, x^p)$ . Unlike many zero-shot TTS models that compress all the information in the speech into a fixed-size vector, we introduce a text encoder  $f_{enc}$ , which takes the text  $c$  and the random segment  $x^p$  as a speech prompt to generate a speaker-conditional text representation. We refer to our text encoder as a speech-prompted text encoder. The output of this text-encoder is then mapped to mel-spectrograms using a flow-matching generative model, and finally to waveform using an audio vocoder.
Speech-Prompted Text Encoder: As shown in Fig. 1a, we input the mel-spectrogram of a random segment  $x^{p}$  as a speech prompt along with the text input  $c$ . We then project both to the same
(a) An overview of P-Flow
(b) Speech-prompted Text Encoder
Figure 1: The overall diagram of P-Flow. P-Flow is a model composed of a speech-prompted text encoder, which outputs a representation containing speaker information from the speech prompt, and a flow matching decoder generates high-quality speech significantly faster than real-time.
dimensions to use as inputs to the text encoder. The role of the speech-prompted text encoder is to generate a speaker-conditional text representation  $h_c = f_{enc}(x^p, c)$  using the speaker information extracted from the prompt  $x^p$ . Similar to large-scale codec language models, we employ a non-autoregressive transformer architecture that can attend to speech prompts at arbitrary text positions.
To train the speech-prompted text encoder to effectively extract speaker information from the speech prompt, we use an encoder loss that directly minimizes the distance between the text encoder representation and the mel-spectrogram. In addition to its original purpose in Grad-TTS [30] for reducing sampling steps in a diffusion-based single-speaker TTS model, it also encourages the encoder to incorporate speaker-related details into the generated text representation.
As speech-prompted encoder output  $h_c$  and the mel-spectrogram  $x$  have different lengths, we align the text encoder output with the mel-spectrogram using the monotonic alignment search (MAS) algorithm proposed in Glow-TTS [21]. By applying MAS, we derive an alignment  $A = MAS(h_c, x)$  that minimizes the overall L2 distance between the aligned text encoder output and the mel-spectrogram. Based on the alignment  $A$ , we determine the duration  $d$  for each text token and expand the encoder output  $h_c$  by duplicating encoder representations according to the duration of each text token. This alignment process results in text encoder output  $h$  that aligns with mel-spectrogram  $x$ . The fairly straightforward reconstruction loss is written as  $L_{enc} = MSE(h, x)$ .
In practice, even though the model is not given the exact positioning of  $x^p$  within  $x$  during training, we found the model to still collapse to a trivial copy-pasting of  $x^p$ . To avoid this, we simply mask out the reconstruction loss for the segment corresponding to  $x^p$ . Despite this, the final model is still capable of inferring a continuous mel-spectrogram sequence. We define the masked encoder loss  $L_{enc}^p$  by using  $m^p$  for the random segment  $x^p$  in the mel-spectrogram  $x$ :
$$
L _ {e n c} ^ {p} = M S E \left(h \cdot m ^ {p}, x \cdot m ^ {p}\right) \tag {1}
$$
By minimizing this loss, the encoder is trained to extract speaker information as much as possible from the speech prompt in order to generate an aligned output  $h$  that closely resembles the given speech  $x$ , which results in enhancing in-context capabilities for speaker adaptation.
Flow Matching Decoders: To perform high-quality and fast zero-shot TTS, we use a flow-matching generative model as the decoder for modeling the probability distribution  $p(x|c,x^p) = p(x|h)$ . Our flow-matching decoder models the conditional vector field  $v_{t}(\cdot |h)$  of Continuous Normalizing Flows (CNF), which represents the conditional mapping from standard normal distribution to data distribution. The decoder is trained using a flow-matching loss  $L_{cfm}^{p}$  that also applies a mask  $m^p$  for the random segment as in  $L_{enc}^{p}$ . More details will be provided in Section 3.2.
Duration Predictor: To reproduce text-token durations during inference where MAS is unavailable, we use a duration predictor trained in a manner similar to [21]. We use the hidden representation
of the speech-prompted text encoder as its input without additional speaker conditioning, given this representation already contains speaker information. It is trained simultaneously with the rest of the model with detached inputs to avoid affecting the text-encoder training. The duration predictor estimates the log-scale duration  $\log \hat{d}$  for each text token, and the training objective for the duration predictor  $L_{dur}$  is to minimize the mean squared error with respect to log-scale duration  $\log d$  obtained via MAS during training.
The overall training loss for P-Flow is  $L = L_{enc}^{p} + L_{cfm}^{p} + L_{dur}$ . Zero-shot inference uses a random chunk from a reference sample as the speech prompt along with the desired transcript as inputs to the speech-prompted text encoder to obtain  $h_c$ . We do not provide a transcript corresponding to the speech prompt. Then, using the predicted duration  $\hat{d}$  from the duration predictor, we expand the encoder output  $h_c$  to obtain  $h$  and generate the personalized speech using the flow-matching decoder.
3.2 Flow Matching Decoder
We use Flow Matching to model the mel-spectrogram decoder task's conditional distribution:  $p(x|h)$ . We first provide a brief overview of flow matching, followed by describing our sampling procedure and additional qualitative improvements through a guidance-related technique.
Flow Matching Overview: Flow Matching [23, 35, 24] is a method for fitting to the time-dependent probability path between our data density  $p_1(x)$  and our simpler sampling density  $p_0(x)$  (assumed to be the standard normal). It is closely related to Continuous Normalizing Flows, but is trained much more efficiently in a simulation-free fashion, much like the typical setup for Diffusion and Score Matching models [13, 16, 34]. We adopt Conditional Flow Matching as specified in [23] as their formulation encourages simpler and often straighter
Figure 2: Conditional Flow Matching [23] can be supervised with linear conditional trajectories during training when mapping between the complex  $p_{\mathrm{data}}$  and the sampling Gaussian, leading to simpler marginal trajectories during inference. Simpler trajectories require fewer function evaluations.
trajectories between source and target distributions. Simpler trajectories allow for test-time sampling in fewer steps without the need for additional distillation. We will ignore the conditional variable  $h$  for notational simplicity in this overview.
Following Lipman et al. [23], we define the flow  $\phi :[0,1]\times \mathbb{R}^d\to \mathbb{R}^d$  as the mapping between our two density functions using the following ODE:
$$
\frac {d}{d t} \phi_ {t} (x) = v _ {t} \left(\phi_ {t} (x)\right); \quad \phi_ {0} (x) = x \tag {2}
$$
Here,  $v_{t}(x)$  is the time-dependent vector field and specifying the trajectory of the probability flow through time.  $v_{t}(x)$  is also our learnable component, henceforth denoted as  $v_{t}(\phi(x); \theta)$ . To sample from the distribution, we sample from the sampling distribution  $p_{0}$  as our initial condition at  $t = 0$  and solve the ODE in Eq. (2). Notably, the formulation [23] encourages straighter trajectories, ultimately allowing us to cheaply approximate the ODE solution with 10 Euler steps with minimal loss in quality.
It turns out that determining the marginal flow  $\phi_t(x)$  is difficult in practice. Lipman thus formulates it as marginalizing over multiple conditional flows  $\phi_{t,x_1}(x)$  as follows:
$$
\phi_ {t, x _ {1}} (x) = \sigma_ {t} (x _ {1}) x + \mu_ {t} (x _ {1}) \tag {3}
$$
Here,  $\sigma_t(x_1)$  and  $\mu_t(x_1)$  are time-conditional affine transformations for parameterization the transformation between Gaussian distributions  $p_1$  and  $p_0$ . Finally, let  $q(x_1)$  be the true but likely non-Gaussian distribution over our data. We define  $p_1$  as a mixture-of-Gaussian approximation of  $q$  by perturbing individual samples with small amounts of white noise with  $\sigma_{\mathrm{min}}$  (empirically set to 0.01). We can specify our trajectories without complications from stochasticity as in SDE formulations.
Taking advantage of this, Lipman et al. recommend simple linear trajectories, yielding the following parameterization for  $\phi_t$ :
$$
\mu_ {t} (x) = t x _ {1}, \sigma_ {t} (x) = 1 - (1 - \sigma_ {\min }) t \tag {4}
$$
Training the vector field is performed using the conditional flow matching objective function:
$$
L _ {C F M} (\theta) = \mathbb {E} _ {t \sim U [ 0, 1 ], x _ {1} \sim q (x _ {1}), x _ {0} \sim p (x _ {0})} \| v _ {t} \left(\phi_ {t, x _ {1}} \left(x _ {0}\right); \theta\right) - \frac {d}{d t} \phi_ {t, x _ {1}} \left(x _ {0}\right) \| ^ {2} \tag {5}
$$
Plugging Eq. (4) in to Eq. (3) and (5), we get our final CFM objective:
$$
L _ {C F M} (\theta) = \mathbb {E} _ {t, q (x _ {1}), p (x _ {0})} \| v _ {t} \left(\phi_ {t, x _ {1}} (x _ {0}); \theta\right) - \left(x _ {1} - \left(1 - \sigma_ {\min }\right) x _ {0}\right) \| ^ {2} \tag {6}
$$
Masked Flow Matching Loss: Recall that our flow matching decoder models the distribution  $p(x|h)$ . Because  $h$  is the output of the text encoder which was provided by the subsegment  $x^p \in x$ , we found it again necessary to mask out the loss for parts of the output corresponding to  $x^p$  to prevent trivial solutions. Let the generic  $v_t(x_t; \theta)$  be parameterized in our setup as  $\hat{v}_\theta(x_t, h, t)$  to account for the conditional. Here,  $t$  is represented with a continuous sinusoidal embedding. This gives us the masked CFM objective:
$$
L _ {C F M} ^ {p} (\theta) = \mathbb {E} _ {t, q (x _ {1}), p (x _ {0})} \| m ^ {p} \cdot \left(\hat {v} _ {\theta} \left(\phi_ {t, x _ {1}} (x _ {0}), h, t\right) - \left(x _ {1} - \left(1 - \sigma_ {\min }\right) x _ {0}\right)\right) \| ^ {2} \tag {7}
$$
Sampling: The conditional flow matching loss marginalizes over conditional vector fields to achieve the marginal vector field, the latter of which is used during sampling. While the linearly interpolated conditional trajectories as specified in Eq. (4) do not guarantee the same degree of straightness in the resulting marginal, we still get something fairly close. Within the context of this work, we found the conditional flow matching formulation to result in simple enough trajectories such that it is sufficient to use the first-order Euler's method with around 10 steps to solve the ODE during inference. Sampling with  $N$  Euler steps is performed with the following recurrence relation:
$$
x _ {0} \sim \mathcal {N} (0, I); x _ {t + \frac {1}{N}} = x _ {t} + \frac {1}{N} \hat {v} _ {\theta} \left(x _ {t}, h, t\right) \tag {8}
$$
Guided Sampling: We find that pronunciation clarity can be further enhanced by applying techniques from a classifier-free guidance method [14]. In a related work, GLIDE [29] amplifies their text-conditional sampling trajectory by subtracting the trajectory for an empty text sequence. We employ a similar formulation, guiding our sampling trajectory away from the average feature vector computed from  $h$ , denoted as  $\bar{h}$ .  $\bar{h}$  is computed by averaging the expanded representation  $h$  along the time axis to obtain a fixed-size vector and then duplicated along the time axis. Let  $\gamma$  be our guidance scale. Our guidance-amplified Euler formulation is as follows:
$$
x _ {t + \frac {1}{N}} = x _ {t} + \frac {1}{N} \left(\hat {v} _ {\theta} \left(x _ {t}, h, t\right) + \gamma \left(\hat {v} _ {\theta} \left(x _ {t}, h, t\right) - \hat {v} _ {\theta} \left(x _ {t}, \bar {h}, t\right)\right) \right. \tag {9}
$$
3.3 Model Details