Title: Spherical Flows for Sampling Categorical Data

Abstract:

This work investigates the learning of generative models for discrete sequences by mapping them into a continuous embedding space. While previous methods generally rely on Euclidean space or the probability simplex, we propose operating on the sphere $\mathbb S^{d-1}$. In this geometric setting, the von Mises-Fisher (vMF) distribution provides a natural noise process and allows for a closed-form conditional score. Although the conditional velocity is typically intractable, we leverage the radial symmetry inherent in the vMF density. This symmetry enables the reduction of the continuity equation on $\mathbb S^{d-1}$ to a scalar ordinary differential equation (ODE) based on cosine similarity, where the unique bounded solution defines the velocity.

We demonstrate that both the marginal score and the marginal velocity defined over $(\mathbb S^{d-1})^L$ can be decomposed into posterior-weighted tangent sums. These components differ solely by per-token scalar weights, thereby facilitating both ODE and predictor-corrector (PC) sampling strategies. The posterior distribution serves as the sole learned component, optimized via a cross-entropy loss. Our experimental evaluation contrasts the proposed vMF path with geodesic and Euclidean alternatives. The results indicate that combining vMF with PC sampling yields significant performance improvements in tasks involving language modeling and Sudoku.

Source: arXiv Generated at: 2026-06-03 00:00:00 UTC