Optimality of FSQ Tokens for Continuous Diffusion for Categorical Data with Application to Text-to-Speech

Continuous diffusion for categorical data is a framework belonging to the diffusion family and aiming at generating discrete data. The scientific interest to such models has been constantly increasing these days because researchers try to achieve a challenging goal of finding reasonable alternatives to autoregressive large language models. In this paper, we study the properties of the structure of the latent space corresponding to discrete tokens expressed in terms of Kullback-Leibler divergence on diffusion path measures and accuracy of the correct token prediction by the optimally trained diffusion model. We find that FSQ tokenization scheme has the latent space structure with the properties that make it best suited for continuous diffusion for categorical data as verified through rigorous theoretical analysis and numerical experiments. To validate our findings in real-life scenario, we train several text-to-speech diffusion models having speech tokens as intermediate acoustic features, and show that the one based on FSQ tokens indeed performs the best, and, moreover, it outperforms its strong LLM-based counterpart, at the same time being significantly smaller and faster.

The paper studies a particular type of diffusion models capable of generating discrete data such as text or speech tokens. Currently, this type of models is not widely used by the practitioners because it heavily relies on the geometry of the latent space associated with it, and optimizing this geometry is challenging. In our paper we apply stochastic calculus to study the properties of the latent space making it optimal for such diffusion models, and conclude that the specific tokenization scheme called FSQ has certain optimality properties. We apply our findings to a text-to-speech algorithm and show that it is possible to replace LLM with significantly smaller diffusion model without loss of quality, which means that properly designed diffusion models can rival LLMs.