Accelerating Codec-based Speech Synthesis with Multi-Token Prediction and Speculative Decoding

TL;DR

The core idea is to predict multiple tokens per inference step of the AR module using multiple prediction heads, resulting in a linear reduction in synthesis time as the number of heads increases, and a novel speculative decoding technique that utilises a Viterbi-based algorithm to select the optimal sequence of generated tokens at each decoding step.

Abstract

The goal of this paper is to accelerate codec-based speech synthesis systems with minimum sacrifice to speech quality. We propose an enhanced inference method that allows for flexible trade-offs between speed and quality during inference without requiring additional training. Our core idea is to predict multiple tokens per inference step of the AR module using multiple prediction heads, resulting in a linear reduction in synthesis time as the number of heads increases. Furthermore, we introduce a novel speculative decoding technique that utilises a Viterbi-based algorithm to select the optimal sequence of generated tokens at each decoding step. In our experiments, we demonstrate that the time required to predict each token is reduced by a factor of 4 to 5 compared to baseline models, with minimal quality trade-off or even improvement in terms of speech intelligibility. Audio samples are available at: multpletokensprediction.github.io/multipletokensprediction.github.io/.

Figures (6)

• Figure 1: AR module with $n = 3$ addition heads. Multiple heads are attached into current model that are simultaneously optimized to predict few future tokens given current state. Note that all the heads are trained once but can be flexibly ignored during inference for faster inference speed. During inference, few future tokens are predicted given current state per inference step. The generated tokens are concatenated into input sequence for the next inference step, respectively. Speculative decoding is also applied to enhance to quality of output tokens sequences.
• Figure 2: Viterbi-based Speculative Decoding is illustrated as follows: (1) Multiple prediction heads generate several distributions per timestep simultaneously. (2) To optimize memory and computational efficiency, the dimensions of the transition matrix and state probabilities are reduced by selecting only the necessary rows and columns. (3) The best sequence is determined using Speculative Decoding, as described in Algorithm \ref{['algo:speculative_decoding']}. The transition matrix computation for LibriTTS is completed in just 3 minutes. Additionally, $top_k$ sampling is employed to preserve diversity.
• Figure 3: TPT$\downarrow$ over num. of heads
• Figure 4: WER$\downarrow$ over num. of heads
• Figure 5: UTMOS$\uparrow$ over num. of heads