Spherical Flows for Sampling Categorical Data

2026-05-06

AI Summary

This work addresses the inefficiency of existing approaches to discrete sequence generation, which lack effective continuous embeddings and sampling mechanisms. The authors propose a generative model formulated on the spherical space $\mathbb{S}^{d-1}$, uniquely integrating spherical flows with the von Mises–Fisher (vMF) distribution. By leveraging the radial symmetry of the vMF distribution, they derive a closed-form conditional score function and an analytically tractable velocity field. Furthermore, they exploit cosine similarity to reduce the continuity equation to a scalar ordinary differential equation (ODE), enabling efficient ODE-based and predictor–corrector sampling solely from the learned posterior distribution. Experimental results demonstrate that the proposed method significantly outperforms baselines operating in geodesic and Euclidean spaces on Sudoku solving and language modeling tasks.

Abstract

We study the problem of learning generative models for discrete sequences in a continuous embedding space. Whereas prior approaches typically operate in Euclidean space or on the probability simplex, we instead work on the sphere $\mathbb S^{d-1}$. There the von Mises-Fisher (vMF) distribution induces a natural noise process and admits a closed-form conditional score. The conditional velocity is in general intractable. Exploiting the radial symmetry of the vMF density we reduce the continuity equation on $\mathbb S^{d-1}$ to a scalar ODE in the cosine similarity, whose unique bounded solution determines the velocity. The marginal velocity and marginal score on $(\mathbb S^{d-1})^L$ both decompose into posterior-weighted tangent sums that differ only by per-token scalar weights. This gives access to both ODE and predictor-corrector (PC) sampling. The posterior is the only learned object, trained by a cross-entropy loss. Experiments compare the vMF path against geodesic and Euclidean alternatives. The combination of vMF and PC sampling significantly improves results on Sudoku and language modeling.

Problem

• categorical data
• generative modeling
• spherical embedding
• discrete sequences
• sampling

Innovation

• spherical flows
• von Mises-Fisher distribution
• conditional score
• predictor-corrector sampling
• discrete generative modeling