Spherical Flows for Sampling Categorical Data

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.

• The paper introduces a principled vMF-based generative flow operating on spherical manifolds to sample categorical data efficiently.
• It reduces complex high-dimensional PDEs to a tractable scalar ODE, enabling stable per-token score and velocity computations.
• Experimental results on Sudoku and LM1B demonstrate superior performance and stability relative to Euclidean and diffusion-based methods.

Spherical Flows for Sampling Categorical Data: An Expert Analysis

Introduction and Motivation

The paper "Spherical Flows for Sampling Categorical Data" (2605.05629) addresses a crucial problem in generative modeling: efficiently learning and sampling from discrete sequence distributions using continuous-time and continuous-space generative models. Classical flow-based and diffusion approaches for discrete data have primarily operated in Euclidean embedding spaces or directly on the probability simplex. This work departs from these paradigms by instead working on the product of unit spheres, $\mathbb{S}^{d-1}$, leveraging the von Mises-Fisher (vMF) distribution—a canonical direction distribution on the sphere—as the basis for both noising and denoising processes.

The motivation is to construct a theoretically principled and practically advantageous generative flow for categorical data that respects the geometry enforced by the common practice of normalization in token embeddings. This approach enables a tractable and geometrically aligned noise and reverse process, crucially yielding closed-form scores and amenable velocities for ODE and predictor-corrector sampling.

Theoretical Contributions

Geometric Formulation on the Sphere

Instead of modeling token sequences in Euclidean space—where token embeddings are typically normalized but noise processes may not preserve this constraint—the authors propose to view the data as points on products of spheres, $(\mathbb{S}^{d-1})^L$, where $L$ is sequence length. This ensures that the forward (noising) and reverse (denoising) processes remain on the manifold to which embedding vectors are constrained.

Conditional Paths with von Mises-Fisher Dynamics

The central theoretical innovation is defining a family of diffusion/flow paths parameterized by the vMF distribution. For a given token embedding $w \in \mathbb{S}^{d-1}$, the forward process at time $t$ is a vMF with mean $w$ and an increasing concentration parameter $\kappa_t$, thus interpolating between uniform noise on the sphere ($\kappa=0$) and a Dirac at $w$ ($\kappa\to\infty$).

Due to the radial symmetry of vMF, the paper shows that the continuity equation governing the measure flow on the sphere can be reduced to a one-dimensional ODE in terms of the cosine similarity $(\mathbb{S}^{d-1})^L$0. This ODE admits a unique bounded solution, which determines the velocity field necessary for sampling and learning.

The authors give an explicit, numerically stable flux-based evaluation method for the per-token velocities and show that both ODE and predictor-corrector (Langevin) sampling are available from a single learned posterior object.

Posterior Learning and Sufficient Statistics

Both the velocity and the score of the marginal densities across the product manifold $(\mathbb{S}^{d-1})^L$1 decompose into sums over vocabulary tokens, weighted by the learned posterior $(\mathbb{S}^{d-1})^L$2 at each position. This observation yields a streamlined training setup: a single cross-entropy loss at every time and token suffices, as the model is required only to learn the per-position posterior distribution over tokens, from which all necessary quantities for sampling can be constructed.

Theoretically, the proposition that the cross-entropy loss is minimized exactly at the true marginal posterior ensures statistical efficiency and optimality in the learned model, paralleling the general score-matching and flow-matching frameworks.

Comparison of Spherical Paths

The paper provides an asymptotic analysis demonstrating that geodesic interpolation (slerp) and vMF-based paths differ in the rate at which signal (i.e., cosine similarity between clean and noisy samples) increases with the schedule parameter and embedding dimension. In particular, for high-dimensional embeddings, vMF interpolants require the schedule to scale with dimension to maintain informative regimes, while slerp paths are less informative at moderate values of $(\mathbb{S}^{d-1})^L$3 in high $(\mathbb{S}^{d-1})^L$4.

Experimental Results

Tasks and Baselines

The method is validated on two challenging discrete tasks:

• Sudoku-Extreme (structured combinatorial reasoning)
• LM1B language modeling (natural language generation at scale)

Competing approaches include discrete-state masked diffusion methods, Euclidean-space continuous flows (VP/VE), and geodesic (slerp) spherical flows. Predictor-corrector (PC) sampling and ODE-only (Euler) sampling are both benchmarked, using DiT-style transformer backbones for all approaches to control for architecture variation.

Numerical Outcomes

Sudoku Generation

• The vMF flow with predictor-corrector sampling achieves a validity rate of 78.2% (non-time-conditioned) and 74.5% (time-conditioned), outperforming masked diffusion ($(\mathbb{S}^{d-1})^L$5), geodesic interpolation (52.2\%), and Euclidean paths.
• Predictor-corrector sampling yields substantial improvements for vMF and VP paths, but not for VE, supporting the claim that bounded score magnitudes (as in vMF) are advantageous for stable and effective Langevin correctors.

LM1B Language Modeling

• The vMF flow sampled with predictor-corrector achieves perplexity as low as 48.5 at entropy $(\mathbb{S}^{d-1})^L$6 (time-conditioned), corresponding to entropy-matched competitive or superior performance to recent discrete and continuous diffusion methods (e.g., DFM Diagonal at PPL $(\mathbb{S}^{d-1})^L$7 with $(\mathbb{S}^{d-1})^L$8).
• ODE sampling only achieves PPL $(\mathbb{S}^{d-1})^L$9 (vMF, tc), highlighting the necessity of PC sampling for harnessing the model's full capacity on this geometry.
• VP and VE baselines remain significantly behind vMF under PC, especially as entropy thresholds decrease.

Empirical Insights

• Posterior-weighted tangent projection structures (enabled by vMF geometry) are crucial for tractable, effective predictor-corrector sampling.
• Euclidean methods exhibit instability or degenerate samples at extreme entropy due to divergent scores, while vMF (being globally regularized by the sphere) avoids this pathology.
• Learning a time-varying noise schedule improves learning and inference convergence, consistent with prior continuous diffusion research [dieleman2022continuous].

Implications and Future Directions

Practical Implications

The formulation of generative modeling with vMF paths on spherical manifolds directly aligns model geometry with data encoding practices in contemporary transformers, likely facilitating better optimization and generalization. The approach particularly benefits cases where inference efficiency (few steps, high quality) is needed and when embedding norms are preserved—features highly relevant in scaled LLMs.

Theoretical Implications

Reducing the high-dimensional PDE over the sphere to a computationally feasible scalar ODE establishes a new technique for manifold-based measure flows, potentially extensible to other symmetries and manifolds common in representation learning. The result foregrounds the integration of geometric measure theory, information geometry, and machine learning.

Relation to Prior Work and Distinctions

• Diffusion in Euclidean Embeddings: Previous approaches [dieleman2022continuous, gulrajani2023likelihood] add Gaussian noise in $L$0, resulting in off-manifold states and requiring ad-hoc solutions for normalization and decoding.
• Simplex- or Assignment Manifold Flows: Other contemporaneous work, such as "Generative Modeling of Discrete Joint Distributions by E-Geodesic Flow Matching on Assignment Manifolds" (Boll et al., 2024) and Dirichlet/assignment flows, targets probabilistic simplex structures, whereas this paper operates directly in normalized embedding spaces.
• Score and Velocity Tractability: The availability of closed-form (or efficiently tabled) scores and velocities on the sphere, not reliant on heavy numerical approximation or discretization, distinguishes this approach and underpins its empirical strengths.

Future Research Prospects

Key avenues include:

• Scaling to Longer Sequences and Larger Vocabularies: Exploiting the efficiency of the per-token decomposition and the stability of vMF flows may facilitate further scaling.
• Hybrid Geometric Flows: Analogous constructions on other homogeneous spaces (tori, hyperbolic spaces) may yield further improvements for domains with intrinsic geometric requirements.
• Optimal Schedules and Adaptive Inference: Joint learning of schedule and flow parameters, as well as adaptive sampling strategies, could close the remaining perplexity/performance gaps.

Conclusion

This work establishes the von Mises-Fisher flow on $L$1 as a tractable, geometrically matched, and empirically superior method for continuous-time generative modeling of categorical data. By leveraging the sphere's geometry, the approach decouples the learning of posteriors from the numerical integration complexities of high-dimensional manifolds, ensuring both statistical efficiency and practical effectiveness. The theoretical reductions and empirical results strongly suggest that geometric congruence between model, data embedding, and generative process is a fruitful principle for future advances in discrete data generation within continuous frameworks.