WaveRNN Architecture for Real-Time TTS on Mobile CPUs
WaveRNN addresses the computational bottlenecks of autoregressive neural audio generation to achieve high-fidelity text-to-speech (TTS) on consumer-grade mobile CPUs. By replacing complex dilated convolutional networks with a streamlined recurrent neural network and combining dual-softmax sample prediction, structured block sparsity, and subscale generation, WaveRNN dramatically cuts memory bandwidth demands and floating-point operations. These targeted architectural modifications allow standard mobile hardware to generate 24 kHz, 16-bit audio in real time.
Dual-Split Softmax Output
Autoregressive audio synthesis requires modeling 16-bit scalar audio samples, which corresponds to 65,536 distinct quantization levels. Predicting a categorical distribution over this full space requires an enormous 65,536-way softmax layer, creating severe memory and compute bottlenecks.
WaveRNN replaces this single massive output layer with a split categorical prediction:
- The 16-bit sample is divided into two 8-bit bytes: a coarse byte (high-order bits) and a fine byte (low-order bits).
- The network uses two consecutive 256-way softmax distributions. The model predicts the coarse byte first, and then conditions the prediction of the fine byte on both the hidden state and the predicted coarse byte.
- This reduces the output computational complexity from \(\mathcal{O}(2^{16})\) down to \(\mathcal{O}(2^8 + 2^8) = \mathcal{O}(512)\), removing the primary output-layer bottleneck on constrained processors.
Structured Block Sparsity
Recurrent neural networks require reloading weight matrices into cache at every time step. For sequential audio generation at 24 kHz, repeatedly moving large matrices between memory and CPU registers limits throughput due to memory bandwidth constraints.
To eliminate this memory bandwidth bottleneck:
- WaveRNN implements non-uniform structured weight pruning, inducing block sparsity (such as 16x1 or 4x4 blocks of weights).
- Pruning up to 96% of the recurrent parameters reduces the number of operations and memory footprints without degrading audio fidelity.
- Because the pruning is structured in discrete blocks rather than individual random weights, the non-zero parameters align with SIMD (Single Instruction, Multiple Data) execution units, such as ARM NEON registers on mobile CPUs. This enables direct hardware vectorization without the latency overhead of sparse-matrix indexing.
Subscale Generation and Parallelism
Standard autoregressive generation advances one sample at a time, preventing multi-core utilization. WaveRNN introduces subscale generation to introduce data-level parallelism into the generation loop.
- The model splits a target audio sequence of length \(T\) into \(B\) interleaved sub-sequences, each of length \(T / B\).
- The recurrent cell maintains \(B\) parallel hidden streams, advancing all streams collectively step-by-step.
- Instead of running \(T\) individual sequential iterations, the CPU processes \(T / B\) steps over a batch size of \(B\).
- This allows mobile multi-core architectures to run multiple sample updates in parallel, maximizing instruction pipelining and preventing pipeline stalls inherent to purely scalar step-by-step recurrence.
Decoupled Conditioning Hierarchy
High-fidelity TTS requires conditioning the waveform generator on auxiliary features such as mel-spectrograms or linguistic alignments. In earlier systems, the dense autoregressive loop frequently re-encoded these conditional features at every sample.
WaveRNN separates the network into two distinct functional hierarchies:
- Conditioning Network: A non-autoregressive feed-forward or convolutional subnetwork that processes low-frequency conditioning features at frame rates (typically 200 Hz). It computes conditioning vectors up front.
- Autoregressive Core: A minimal, lightweight gated recurrent unit (GRU) running at 24 kHz that consumes the pre-computed conditioning tensors via simple linear projections.
By offloading heavy feature transformations to the frame-rate network, the sample-level recurrent core is left with only the bare minimum operations necessary to guarantee temporal coherence, preserving critical CPU cycles during generation.