PaperlayerView exact revision on arXiv

Source-linked AI summary

Language Modeling with Hyperspherical Flows

Justin Deschenaux, Caglar Gulcehre

arXiv:2605.11125v3cs.LG

TL;DR

Discrete diffusion generates in parallel but its factorized sampling is less expressive than autoregressive modeling, while prior flow language models incur one-hot training costs. S-FLM uses a hyperspherical embedding flow, and at standard temperature it closes the gap to masked diffusion, although a substantial gap remains at low temperature and to autoregressive decoding.

  • Problem

    Discrete diffusion’s factorized token sampling is less expressive than AR, while prior flow language models use costly one-hot representations and Gaussian noise without clear semantic interpretation.

    Source: §1 · p. 2

  • Method

    S-FLM learns a velocity field over token embeddings on the hypersphere and uses top-k velocity decoding, with k = 1 corresponding to greedy decoding.

    Source: §1 · p. 3, §3.1 · p. 5

  • Results

    At T = 1, S-FLM closes the gap to MDLM and Duo with top-1 decoding across NFE budgets and beats both with exact velocity at NFE ≤16.

    Source: §1 · p. 3

  • Takeaways & Limitations

    S-FLM avoids materializing vocabulary-sized one-hot vectors and is competitive with masked diffusion on GSM8K at standard-temperature sampling.

    Source: §1 · p. 3

    On TinyGSM, S-FLM reduces but does not eliminate the gap to prior FLMs, and a substantial gap to autoregressive decoding remains.

    Source: §7 · p. 10

Abstract

from arXiv · show

Discrete Diffusion Language Models progressed rapidly as an alternative to autoregressive (AR) models, motivated by their parallel generation abilities. However, for tractability, discrete diffusion models sample from a factorized distribution, which is less expressive than AR. Recent Flow Language Models (FLMs) apply continuous flows to language, transporting noise to data with a deterministic ODE that avoids factorized sampling. FLMs operate on one-hot vectors whose dimension scales with the vocabulary size, making FLMs costly to train. Moreover, since all distinct one-hot embeddings are equidistant in $\ell_2$, adding Gaussian noise does not have a clear semantic interpretation (unlike images, where Gaussian noise progressively degrades structure). We introduce $\mathbb{S}$-FLM, a latent FLM in the hypersphere. $\mathbb{S}$-FLM generates sequences by rotating vectors in $\mathbb{S}^{d-1}$ along a velocity field learned with cross-entropy, avoiding the overhead of materializing one-hot vectors. Previous FLMs match AR in Generative Perplexity (Gen.\ PPL), but samples with high likelihood are not necessarily correct in verifiable domains such as math and code. $\mathbb{S}$-FLM substantially improves continuous flow language models on large-vocabulary reasoning and closes the gap to masked diffusion under standard-temperature sampling ($T=1$), while a gap remains under optimized low-temperature ($T=0.1$) decoding.

1 Introduction

Discrete diffusion supports parallel generation but samples from factorized token distributions, while flow language models use continuous transport to update positions jointly. S-FLM moves token embeddings onto a hypersphere to avoid one-hot costs, and on GSM8K it closes the gap to masked diffusion at T = 1, though a low-temperature gap remains.

Motivation

  • Discrete diffusion models generate in parallel with bidirectional context, but factorized marginals make them less expressive than autoregressive models.AR models use chain-rule factorization and scale to large sizes, but require one forward pass per generated token and causal attention can hurt reasoning tasks.Source: §1 · p. 2

Background

  • Continuous flows transport noisy samples to data through an ODE, jointly updating positions and avoiding discrete diffusion’s factorized sampling.Recent FLMs apply this approach to language using one-hot vectors, Gaussian noise, and a cross-entropy-trained denoiser.Source: §1 · p. 2

Motivation

  • Prior FLMs match AR and discrete diffusion in Gen. PPL, but one-hot representations are costly to train at large vocabulary sizes and Gaussian noise has unclear meaning on them.With 100k–200k-token vocabularies, one-hot vectors exceed 100k dimensions and noisy inputs require multiplication by the embedding matrix rather than a lookup.Source: §1 · p. 2

Contributions

  • At T = 1, S-FLM closes the gap to MDLM and Duo with top-1 velocity decoding, and exact velocity beats both at NFE ≤16.At T = 0.1, MDLM and Duo reach 33–36%, while S-FLM remains around 18%.Source: §1 · p. 3
  • S-FLM is a Riemannian flow over embeddings that operates in d dimensions without materializing |V|-dimensional one-hot vectors.Its S-arch backbone has normalized activations; the paper reports that aligning activations with the input improves sample quality on GSM8K and OWT.Source: §1 · p. 3

2 Background

The background establishes language-model notation and contrasts autoregressive sequence factorization with continuous transport models. It then introduces flow matching and the spherical geometry and training formulation used by the paper.

2 Background

  • The notation defines V as a finite vocabulary and describes AR likelihoods as chain-rule products that generate tokens sequentially without conditioning on future tokens.Sequences contain L tokens, and token embeddings are rows of a lookup table E in R^{|V|×d}.Source: §2 · p. 3

2.1 Flow Generative Modeling

  • Continuous normalizing flows transport noise to data by integrating a neural velocity field as an ODE, producing intermediate distributions along the path.Flow matching trains the field against conditional velocities because the posterior given a noisy sample is generally intractable; its minimizer is the marginal velocity.Source: §2.1 · pp. 3–4
  • Conditional flow matching uses a noise schedule and linear interpolation between noise and data, with the conditional velocity serving as its training target.The conditional velocity depends on the data sample and noisy example, and minimizing the loss recovers the marginal velocity.Source: §2.1 · pp. 3–4

2.2 Geometry of the Hypersphere

  • On the unit hypersphere, geodesic distance is angular, tangent spaces contain vectors orthogonal to their base point, and the logarithmic map returns a tangent vector toward another point.SLERP follows a geodesic, while a uniform spherical sample can be obtained by normalizing an isotropic Gaussian vector.Source: §2.2 · p. 4

2.3 Flow Matching on the Hypersphere

  • Riemannian Flow Matching extends flow matching to manifolds such as the sphere, where SLERP connects noise and data and defines the conditional flow.The conditional velocity is obtained by differentiating that flow, and the minimizer of the spherical conditional-flow-matching loss is the marginal velocity.Source: §2.3 · p. 4

3 Hyperspherical Flow Language Models

S-FLM represents sequences with unit-norm token embeddings on a hypersphere and trains a cross-entropy denoiser whose posterior induces a flow velocity. Its noise schedule and sphere-preserving architecture shape how that flow is trained and sampled.

3 Hyperspherical Flow Language Models

  • S-FLM maps tokens to unit-norm sphere embeddings and trains a cross-entropy denoiser whose posterior yields a marginal velocity that is integrated for sampling.The denoiser’s cross-entropy objective pushes token embeddings apart rather than allowing them to collapse.Source: §3 · pp. 4–5

3.1 Sampling

  • Sampling uses an exact posterior-weighted velocity, a one-sample stochastic alternative, or a top-k posterior truncation, with k = 1 corresponding to top-1 decoding.The exact velocity is a finite sum over token representations; stochastic decoding replaces that sum with a posterior sample.Source: §3.1 · p. 5

3.2 Noise Schedule

  • Truncating training to noise levels in [0, a] avoids levels where the posterior is near one-hot; a vocabulary- and dimension-based bound performs similarly to the best GSM8K grid-search value.An adaptive schedule also allocates more samples to noise levels with greater loss change, without slowing training.Source: §3.2 · pp. 5–6

3.3 Hyperspherical Architecture

  • S-arch keeps intermediate activations on the sphere with rotational attention and MLP layers and normalized interpolation, achieving better results with slightly fewer parameters than standard DiT.Its per-dimension interpolation gate depends on noise level, and the normalized interpolation approximates SLERP for small γ.Source: §3.3 · p. 6

4 Experiments

The experiments evaluate S-FLM on verifiable Sudoku and math tasks as well as OpenWebText, since likelihood alone may not indicate sequence-level correctness. Results show strong gains over prior continuous methods on GSM8K and stable low-NFE language-modeling frontiers, with remaining gaps to discrete diffusion in low-temperature reasoning.

4 Experiments

  • Because Gen. PPL may not reflect sequence-level correctness, the experiments prioritize datasets with ground-truth solutions while also evaluating unconditional OpenWebText modeling.The verifiable tasks are Sudoku and GSM8K.Source: §4 · p. 6

4.1 Reasoning on Sudoku

  • S-FLM and FLM outperform AR and MDLM at every Sudoku difficulty; truncation and the adaptive schedule bring S-FLM near prior best continuous models.S-FLM is reported as best but similar to FLM.Source: §4.1 · p. 7, Table 1 · p. 7

4.2 Reasoning on GSM8K

  • 18.0% accuracy with top-k = 1 closes S-FLM’s gap to MDLM and Duo at T = 1, while MDLM and Duo reach 33–36% at T = 0.1.At T = 1 and NFE ≤16, S-FLM also outperforms MDLM and Duo.Source: §4.2 · p. 8, Table 2 · p. 8

4.3 Language Modeling on OpenWebText

  • At NFE = 1024, S-arch matches prior FLMs in Gen. PPL; at low NFE its frontier remains stable while prior FLMs destabilize, though Duo leads the trade-off.The frontier compares Gen. PPL and average unigram entropy across sampling temperatures.Source: Figure 4 · p. 8, §4.3 · p. 9

5 Related Work

S-FLM differs from earlier language flows in its continuous hyperspherical representation and joint learning of token embeddings and velocity. This design avoids vocabulary-sized dense arrays while situating the method among prior embedding-based flows and hyperspherical representation work.

Continuous diffusion for language modeling

  • Unlike prior manifold flows that generally assume data already lie on the manifold, S-FLM jointly learns embeddings and velocity on the hypersphere and injects noise by rotations.Earlier embedding approaches either train end-to-end with cross-entropy or regress onto pretrained embeddings in two stages.Source: §5 · p. 9

Flow language models (FLMs)

  • S-FLM uses d-dimensional token embeddings instead of prior FLMs’ dense L × |V| arrays, avoiding their vocabulary-size bottleneck and training faster.Fisher-Flow maps one-hot vectors to the sphere but does not scale well to large vocabularies.Source: §5 · p. 9

Representation learning on the hypersphere

  • Hyperspherical representations are linked in prior work to cosine-based embedding comparisons, downstream performance, and stability in variational and autoregressive models.The cited work discusses uniform spread on the sphere, retrieval, latent priors, and normalized activations and weights.Source: §5 · p. 9

6 Conclusion

The conclusion presents S-FLM as a jointly learned Riemannian flow that is competitive on large-vocabulary language modeling. Its reasoning results improve on continuous-flow baselines, but the comparison with discrete diffusion and autoregressive models remains qualified by temperature and task.

6 Conclusion

  • S-FLM learns embeddings and velocity jointly on the hypersphere and is competitive at language modeling with large vocabularies.The S-arch backbone and analysis of truncated and adaptive noise schedules are additional contributions.Source: §6 · p. 9
  • On Sudoku, S-FLM performs similarly to prior flow language models; on GSM8K, top-1 decoding closes the gap to MDLM and Duo at T = 1, but not at T = 0.1.The conclusion also notes a remaining gap between diffusion and autoregressive models on GSM8K.Source: §6 · p. 9
  • On OpenWebText, S-FLM follows prior FLMs’ Gen. PPL/Entropy frontier at high NFE and beats them at low NFE, where their frontier is unstable.The conclusion says avoiding |V|-dimensional one-hot vectors makes S-FLM faster to train; easier scaling is explicitly speculative.Source: §6 · p. 9
  • The truncation analysis gives a principled heuristic for choosing a noise-schedule truncation bound that works well in practice.The bound is derived as a function of vocabulary size and embedding dimension.Source: §6 · p. 9

7 Limitations

The paper identifies performance gaps for continuous diffusion language models, alongside limitations in its sampling heuristic, throughput, training analysis, and experimental replication. Its conclusions therefore remain qualified by both modeling assumptions and training constraints.

7 Limitations

  • Continuous diffusion models, including S-FLM, underperform discrete counterparts; on TinyGSM, S-FLM reduces the gap to prior FLMs but leaves a substantial gap to autoregressive decoding.The authors also report a remaining gap under optimized low-temperature decoding.Source: §7 · p. 10
  • The truncation threshold α⋆(δ) assumes randomly distributed spherical embeddings, so it is a heuristic for learned embeddings that may be more structured.The authors suggest more sophisticated modeling of sampling dynamics might improve the bound.Source: §7 · p. 10
  • The S-arch trains slower than the standard DiT, making improved throughput a stated future direction.Source: §7 · p. 10
  • Only one model was trained per configuration because a TinyGSM training run costs more than $300.The authors say training dynamics should also be studied further.Source: §7 · p. 10

8 Impact Statement

The impact statement describes the work as methodological and notes its small-scale training and limited capabilities. The accompanying sections explain the flow-matching framework, spherical training and sampling, adaptive scheduling, and embedding behavior.

8 Impact Statement

  • The models were trained at small scale on Sudoku, TinyGSM, and OpenWebText, remain far below state-of-the-art autoregressive capabilities, and present standard language-modeling misuse risks.The impact statement characterizes the contribution as methodological.Source: §8 · p. 10

A.1 The Continuity Equation

  • The continuity equation links the velocity field to density evolution, allowing flow matching to target the velocity without computing the flow or its Jacobian.Source: §A.1 · p. 18

A.2 Deriving the Conditional Velocity Fields

  • Conditional flow matching derives tractable regression targets by differentiating a conditional path and expressing its velocity using the current state.The spherical case uses a geodesic flow and its velocity is proportional to the logarithmic map toward the endpoint.Source: §A.2 · pp. 18–19

B.1 Training and Sampling Pseudocode

  • The training and sampling procedures construct noisy latents, update them through successive model evaluations, and decode the final latent as tokens.The deterministic sampler uses one forward pass per step; the stochastic version samples a token and uses its log map.Source: §B.1 · pp. 19–20

B.2 Adaptive Noise Schedule

  • The adaptive schedule allocates more training samples to noise levels where the loss changes fastest, using the loss-profile derivative and smoothing successive schedules with an EMA.The approach is motivated by the observation that very low and very high noise levels leave the denoiser little to learn.Source: §B.2 · p. 20

B.3 Hyperspherical Backbone Architecture

  • The hyperspherical backbone adapts nGPT with bidirectional attention, GELU activations, time-conditioned residual gates, and normalized weight vectors.The gates are constrained to be non-negative, so each update pulls the hidden state toward the block output.Source: §B.3 · pp. 21–22

B.4 Step Size with the Euler Sampler

  • For a linear noise schedule with uniform time steps, the Euler sampler's step size is 1/(N − n).Source: §B.4 · p. 22

B.5 Re-Normalizing Embeddings

  • Normalizing embeddings in the forward pass makes their parameter norms grow during optimization, reducing effective angular learning rates; re-projection prevents that implicit decay, which the authors find beneficial.Source: §B.5 · pp. 22–23

B.6 Training Cost

The section reports training-cost context for Sudoku, TinyGSM, and OpenWebText, and explains that tokenizer choice changes TinyGSM sequence lengths. SmolLM tokenization compresses code better than GPT-2 tokenization.

B.6 Training Cost

  • Sudoku training used one L40S GPU for under two hours and cost less than $2.Source: §B.6 · p. 23
  • TinyGSM training ran for 250k steps on 8 H100s, while OpenWebText training ran for 1M steps on 16 H100s.Both used bfloat16 mixed precision.Source: §B.6 · p. 23

C.1 Sequence Length Using Different Tokenizers

  • On TinyGSM, SmolLM tokenization has a 180-token median sequence length versus 204 tokens for GPT-2, fitting more problems into a 512-token context.The authors attribute the improved compression to SmolLM's code training.Source: Figure 5 · p. 24, §C.1 · p. 24

C.2 Analysis under the Random Codebook Model

The analysis derives a critical noise threshold under a random-codebook model and reports how that threshold varies with embedding dimension and confidence. The Sudoku setup specifies a shared sequence format and training configuration for its models.

C.2 Analysis under the Random Codebook Model

  • The threshold analysis assumes random spherical embeddings and uses concentration bounds for inner products to bound similarity among codewords not selected by simplified generation dynamics.The derivation approximates independently sampled spherical vectors as nearly orthogonal in high dimensions.Source: §C.2 · pp. 24–25
  • The critical threshold α⋆(δ) decreases as embedding dimension increases; lowering δ also lowers the minimum noise level, with Table 4 covering vocabulary size, dimension, and confidence settings.Source: §C.2 · p. 26, Table 4 · p. 26

C.3 Sudoku Setup

  • Sudoku examples use a flattened partial grid with missing digits encoded as 0, followed by a separator and solution; diffusion leaves the grid uncorrupted and excludes its positions from loss.Models use a context length of 180 and a 12-token vocabulary, except MDLM, which adds a mask token.Source: §C.3 · p. 26
  • All Sudoku models train for 20k steps with batch size 256, learning rate 3 × 10^-4, and EMA rate 0.9999.S-FLM uses linear and cos^2 schedules truncated according to the threshold in Table 4.Source: §C.3 · p. 26, Table 4 · p. 26

C.4 Additional Ablations on Sudoku

Truncation improves TinyGSM accuracy, while Sudoku ablations show gains from truncating the noise schedule and skipping embedding renormalization. The section also explains a bootstrap interval procedure for TinyGSM accuracy.

C.4 Additional Ablations on Sudoku

  • Truncating to α⋆(0.1) improves accuracy on Sudoku and TinyGSM; on Sudoku, gains are largest for linear schedules, while untruncated Cosine2 outperforms linear.Source: §C.4 · p. 26
  • Skipping embedding renormalization improves nearly all schedules and difficulties; on hard Cosine2 without truncation, the gap reaches 7.3 points.Source: §C.4 · p. 26, Table 6 · p. 27
  • Normalized-embedding gradients are tangent to the sphere, so gradient steps increase embedding norms, equivalent to learning-rate decay; empirically, this annealing improves Sudoku performance.Source: §C.4 · p. 26

C.5 Bootstrap Confidence Intervals on TinyGSM

  • TinyGSM accuracy error bars use 95% percentile bootstrap confidence intervals from 1,000 resamples of the 1,319 GSM8K test outcomes.Source: §C.5 · p. 27

C.6 Additional Ablations on TinyGSM

  • Accuracy rises from 1.21% at αtrunc = 0 to 7.73% at αtrunc = 0.9; α⋆(0.1) = 0.879 and α⋆(0.01) = 0.869 perform similarly to the best run.The sweep uses a linear schedule and standard DiT backbone trained for 250k steps.Source: §C.6 · p. 27, Table 7 · p. 27

C.7 Additional Sampling Results on TinyGSM

On GSM8K, S-FLM’s performance depends on sampling temperature and velocity decoding, and low-temperature results remain below Duo and MDLM. On OpenWebText, S-FLM traces a Gen. PPL–entropy Pareto frontier across the reported NFE budgets.

C.7 Additional Sampling Results on TinyGSM

  • At T = 0.1, GSM8K accuracy is 36.0% for Duo, about 33% for MDLM, and near 18% for S-FLM.Source: §C.7 · p. 27
  • At T = 1, top-1 reaches 18.0% while k ≥10 and unrestricted decoding plateau near 12%; T = 0.1 improves accuracy by about six points for both backbones.Source: §C.7 · p. 27, Figure 8 · p. 28, Figure 7 · p. 28

C.8 Gen. PPL / Entropy Frontier Protocol

  • The OpenWebText frontier sweeps 15 temperatures and six NFE budgets over 512 length-1024 samples per cell, plotting GPT-2-large Gen. PPL against unigram entropy in a filtered window.Source: §C.8 · p. 29

C.9 Additional Results on OpenWebText

  • S-FLM traces the OpenWebText Gen. PPL–unigram-entropy Pareto frontier at every NFE budget, with its largest margin over MDLM, Duo, and FLM at low NFE.Source: §C.9 · p. 29, Figure 10 · p. 30

D Extended Related Work

  • S-FLM combines continuous flow on a hypersphere with jointly learned token representations, avoiding materialization of dense L × |V| arrays during training.Source: D Extended Related Work · pp. 30–31

Related papers

Loading 2605.11125v3…