Language Modeling with Hyperspherical Flows

In plain words

Most language models write one word after another. A newer family, called diffusion models, starts from random noise over a whole passage and cleans it up in steps, so each step sees the whole passage at once. Most text versions of this work on discrete symbols, and the ones that work on continuous vectors have trailed behind on hard problems such as math.

This paper keeps the vectors but puts them on the surface of a sphere. Each word is a point on that surface, and the direction of a point is what matters, not its length. Generation starts from random points on the sphere and rotates each one toward the point for the word it should become. A network learns, at every step, which words the noisy point is most likely meant to be, and it chooses the direction of the next rotation from those guesses. The word embeddings are learned along with the network, not fixed in advance.

The main claim is that this works much better on a grade-school math benchmark than earlier continuous versions, which were almost unable to solve any problem. Most of the gain comes from three choices: training on only part of the noise range, choosing the training noise levels adaptively from the loss, and letting each step use only the most likely token. Keeping the network's activations on the sphere adds a little more. The model is still behind the best discrete diffusion models and autoregressive models at low sampling temperature.

TL;DR

• S-FLM is a flow over unit-norm token embeddings on the sphere S^(d-1). It denoises by rotation and never builds one-hot vectors. It is trained with unweighted cross-entropy §3 eq. 14.
• Three things carry the result: a truncated and adaptive noise schedule, the S-arch backbone whose activations stay on the sphere, and top-1 velocity decoding §3.2 §3.3 §4.2 tab. 2.
• Author-reported GSM8K accuracy at T=1 with 1024 steps: 1.2% for the plain method, 12.4% with S-arch, 18.0% with top-1 velocity. Prior continuous flow LMs (CANDI, FLM) reach 0.2% and 0.3% tab. 2.
• Use it when you want a continuous, embedding-space flow LM and can afford a per-step sum over the vocabulary. Expect gaps to discrete diffusion at T=0.1, and treat the released sampler defaults with care (see Paper vs code).

Problem

Discrete diffusion language models sample each token from a factorised marginal at each denoising step, which is less expressive than an autoregressive factorisation §1. Earlier continuous flow language models (FLMs) act on one-hot vectors whose size grows with the vocabulary, which makes them costly to train, and Gaussian noise on one-hot vectors has no clear semantic meaning §1 §2. On GSM8K, previous FLMs trained on TinyGSM are below 1% accuracy §1. Generative perplexity, the usual metric, does not show whether a sample is correct, so the paper evaluates on Sudoku and GSM8K, which have ground truth §4.

Method

Overview. S-FLM is a Riemannian continuous normalising flow on (S^(d-1))^L. A sequence of random unit vectors is transported toward the clean token embeddings §3. The denoiser p(x | z_t) is trained with cross-entropy and induces a closed-form velocity field, which is integrated with a geodesic Euler method and decoded by argmax §3 §3.1 alg. 1 alg. 2.

1. Embedding and decoding. Each token v has a row e_v of a learned table E. Its clean endpoint is the unit-norm row e_v / ||e_v|| eq. 13. Decoding at the final time takes the argmax of the posterior, g = argmax_v p(x = v | z) eq. 13 §3.

2. Noise on the sphere. The prior is a uniform point on the sphere, sampled as a normalised Gaussian vector §2.2 alg. 1. The noisy latent is the spherical linear interpolation between the prior and the clean endpoint, with the interpolation level set by the schedule alpha §3 eq. 8 eq. 10. The gradient flows through this SLERP into the embedding table, so the table is learned jointly §3 alg. 1.

3. Training loss. Cross-entropy of the clean token at each position, with no weight on the noise level §3 eq. 14 algo.py. The paper writes the loss as a sum over positions; the code averages over valid tokens in the batch (see Paper vs code).

4. Velocity. Let p(v | z) be the denoiser's posterior at each position. The exact velocity is the posterior-weighted sum of log maps toward every normalised embedding eq. 15. It is an exact marginal velocity only if the posterior is calibrated §3.1.

5. Sampling. Start from a uniform point at each position. At each step, take the velocity from step 4 and move along the geodesic by a step size set by the schedule alg. 2 app. B.4. After the last step, decode by argmax eq. 13.

6. Top-k velocity. Keep only the k largest posterior entries and renormalise before computing the velocity §3.1 samplers.py. Top-1 reduces the velocity to a single log map toward the most likely token, which is a greedy mode-following approximation §3.1. The stochastic variant replaces the sum with one sampled token §3.1 alg. 3.

7. Noise schedule. Training uses a truncated range [0, a] of the schedule and an adaptive density over that range §3.2. The truncation bound comes from a random-codebook model and is given by eq. 17 eq. 17 app. C.2. The adaptive schedule refits every 50 steps from the recent (t, loss) pairs, using the positive part of the derivative of a spline fit to the loss, and smooths the result with an EMA §3.2 alg. 4 app. B.2.

8. S-arch backbone. A transformer whose activations stay on the sphere, after nGPT. Each residual update is a normalised interpolation whose per-dimension gate is time-modulated by a zero-initialised linear map §3.3 app. B.3. Attention uses bias-free projections, RoPE, normalised queries and keys, and softmax scale sqrt(d_k) app. B.3 sphere_arch.py. The MLP is GELU §3.3 app. B.3.

Training pseudocode, following the structure of Algorithm 1 alg. 1:

# One training step  [alg. 1] [§3]
# Input: tokens (B, L); embedding table E (V, d); schedule alpha(t)
# Output: scalar loss

def training_step(tokens, embedding_table, schedule):
    # Clean latent: unit-norm embedding of each token  [§3] [alg. 1]
    clean = normalize(embedding_table[tokens])  # (B, L, d)

    # Prior latent: uniform on the sphere, one per position  [§2.2] [alg. 1]
    prior = normalize(random_normal(clean.shape))  # (B, L, d)

    # Time drawn from [1e-3, 1] with antithetic sampling, then the schedule
    # [§3.2] [code: trainer_base.py#L497-L506]
    t = sample_time(tokens.shape[0])  # (B,)
    alpha = schedule(t)  # (B,), signal level

    # Noisy latent on the geodesic from prior to clean  [§3] [eq. 10]
    noisy = slerp(prior, clean, alpha)  # (B, L, d)

    # Timestep input is sigma = -log(alpha)  [code: trainer_base.py#L521-L522]
    log_probs = log_softmax(denoiser(noisy, -log(alpha)))  # (B, L, V)

    # Unweighted cross-entropy on the clean token  [§3] [eq. 14]
    loss = mean(-gather(log_probs, tokens))  # eq. 14 writes a sum
    return loss

Sampling pseudocode, following Algorithm 2 and the top-k option:

# Deterministic sampler, exact or top-k velocity  [alg. 2] [§3.1]
# Input: denoiser; embedding table E (V, d); steps N; top_k or None
# Output: token ids (B, L)

def sample(denoiser, embedding_table, num_steps, top_k=None):
    # Start from a uniform point on the sphere at each position  [alg. 2]
    latents = normalize(random_normal((B, L, d)))  # (B, L, d)

    # Time grid from pure noise to clean; code time runs 1 -> 0
    # [app. B.4] [code: samplers.py#L635-L640]
    times = linspace(1.0, eps, num_steps + 1)  # (N + 1,)
    unit_table = normalize(embedding_table)  # (V, d)

    for n in range(num_steps):
        # One forward pass gives the posterior over the vocabulary  [§3.1]
        probs = softmax(denoiser(latents, times[n]))  # (B, L, V)
        if top_k is not None:
            # Keep the k largest entries and renormalise  [§3.1]
            probs = top_k_renormalised(probs, top_k)  # (B, L, V)

        # Exact or top-k velocity: posterior-weighted log maps  [eq. 15]
        velocity = log_map_sum(latents, probs, unit_table)  # (B, L, d)

        # Step size from the schedule  [app. B.4]
        step = (alpha(times[n + 1]) - alpha(times[n])) / (
            1 - alpha(times[n]))  # scalar

        # Geodesic Euler step on the sphere  [alg. 2]
        latents = exp_map(latents, step * velocity)  # (B, L, d)

    # Final decode: argmax of the posterior at the last time  [eq. 13]
    final = denoiser(latents, times[num_steps])  # (B, L, V)
    return argmax(final, dim=-1)  # (B, L)

Key equations, in the paper's notation:

\mathbf{z}_t = \mathrm{SLERP}(\mathbf{z}_0, \mathbf{z}_1, \alpha_t) eq. 10

u^\theta_t(\mathbf{z}_t^{(\ell)}) = \frac{\dot{\alpha}_t}{1 - \alpha_t} \sum_{v \in \mathcal{V}} p^\theta_{1|t}(x^{(\ell)} = v \mid \mathbf{z}_t)\, \log_{\mathbf{z}_t^{(\ell)}}(\hat{\mathbf{e}}_v) eq. 15

\alpha^\star(\delta) \approx \frac{2}{\pi} \arcsin\!\left( \sqrt{\frac{2 \log\!\bigl(2(|\mathcal{V}|-1)/\delta\bigr)}{d}} \right) eq. 17

Hyperparameters

Values from the paper's text and appendix. "Not stated" means the paper does not give the value. Code values are the repository's defaults or the released script values, and are labelled as code.

Results

Author-reported. Setup: Sudoku with 180 sampling steps and an 8-layer DiT §4.1; GSM8K with TinyGSM training, SmolLM-135M tokens, 1024 sampling steps for diffusion variants and 512 context for AR §4.2 tab. 2; OpenWebText with 1M training steps §4.3. Error bars on GSM8K are bootstrap confidence intervals over the test set, not over training runs app. C.5.

Sudoku, exact-match accuracy (%), from tab. 1:

GSM8K, accuracy (%) after 250k training steps, from tab. 2. Each S-FLM row adds one component to the row above it:

At T = 0.1, the paper reports MDLM at about 33% and Duo at 36%, while S-FLM with top-1 reaches about 18% §4.2 fig. 3 app. C.7. At T = 1, the exact velocity and the stochastic variant plateau near 12%, and top-1 reaches 18.0% fig. 1 app. C.7.

OpenWebText, generative perplexity against unigram entropy (frontier plot): S-FLM with the S-arch is roughly on par with prior FLMs; Duo is best overall; at NFE 32 the FLM baseline's frontier is unstable fig. 4 §4.3. The paper does not give a single headline number for OpenWebText.

Training cost, from tab. 3 (steps per second and GPU-hours on H100):

Paper vs code

Repository: https://github.com/jdeschena/s-flm at commit 30def79976ca3b1cdcb5fcbfa0165524e44001de. The abstract links it. The owner name matches the first author's site (jdeschena.com), which the README links. The repository has no LICENSE file. The README states an Apache-2.0 license, but no file backs that up, so the code is not quoted here. Only paths and lines are cited.

• Truncation direction, two conventions. The paper's text says the schedule is truncated to [0, a], "i.e., we train only at the higher noise levels" §3.2, with a = alpha*(delta) from eq. 17. For vocabulary about 50k and d = 768 eq. 17 gives 0.121 tab. 4. The released adaptive scripts use exactly this value as the upper bound of the schedule range [code: scripts/train/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.sh#L32] sfm.sh. The released fixed-truncation DiT script instead sets a lower bound of 0.869 and leaves the upper bound at 1 [code: scripts/train/tinygsm/sfm_sphere_dit_truncated_fixed_no_renorm.sh#L28-L29]. That keeps the other end of the schedule. The paper's ablation tab. 7 reports alpha_trunc = 0.869 and 0.879 and calls 0.879 alpha*(0.1) tab. 7 app. C.6. Those values are the lower-bound reading of the fixed script and the complement of eq. 17's 0.121. The table does not state its convention. The paper's text and the adaptive scripts use the upper-bound reading with 0.121. I did not run either reading.

• Sudoku truncation value. The paper text does not name the (|V|, d) pair behind the Sudoku bound. The released Sudoku script sets alpha_max = 0.093, which is the table value for |V| = 12, d = 512, delta = 0.1 [code: scripts/train/sudoku/sfm_truncated_adaptive.sh#L26] tab. 4. So the Sudoku run uses the upper-bound convention.

• Embedding re-projection versus weight re-projection. The paper's "no re-projection" row tab. 6 refers to the embedding table. In code, two switches do different jobs. For S-arch, algo.renormalize_weights = True re-projects every weight matrix after each optimiser step algo.py [code: models/sphere_arch.py#L324-L353], and model.normalize_input_embed = False leaves the embedding table free sphere_arch.py. The released GSM8K S-arch scripts set both, so they keep matrices on the sphere and do not re-project the embeddings. The script names say "no_renorm", which hides the matrix re-projection. For the DiT, renormalize_weights only re-projects the embedding table sphere_dit.py.

• S-arch requires matrix re-projection. The paper says the weight matrices can be re-projected "optionally" app. B.3. The code refuses to run S-arch unless algo.renormalize_weights is True algo.py.

• OWT backbone. The paper says OWT uses the S-arch §4.3 and the README calls the OWT checkpoint "S-arch" README.md. The released OWT training and sampling scripts build small-sphere-dit, the DiT variant [code: scripts/train/owt/sfm.sh#L20] sfm.sh. I cannot tell from the repository which backbone the released OWT checkpoint has.

• Final decoder. Eq. 13 decodes by argmax. The sampler config defaults to ancestral (sampled) decoding sfm.yaml, but every released sampling script overrides it with greedy [code: scripts/sample/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.sh#L41] [code: scripts/sample/sudoku/sfm_truncated_adaptive.sh#L39]. Use greedy to match the reported numbers.

• Sampling steps. The paper uses 1024 steps for GSM8K diffusion baselines and S-FLM tab. 2. The released S-FLM sampling scripts default to 32 steps [code: scripts/sample/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.sh#L12] and the README says the same README.md. The sampler config defaults to 1024 sfm.yaml. Pass the step count explicitly.

• Top-k renormalisation. The paper says top-k "truncates" the posterior to its k largest entries §3.1. The code renormalises those k entries with a log-softmax samplers.py before the velocity [code: samplers.py#L708-L711]. With k = 1 this makes the probability a one-hot, so the sampled and exact velocities are the same. The README's "top-1, sampled" recipe README.md therefore equals TOPK_VELOCITY=1 alone.

• Loss normalisation. Eq. 14 is a sum over positions eq. 14. The base trainer averages the per-token loss over valid tokens in the batch [code: trainer_base.py#L447-L456]. Adam's update is largely unaffected by a constant loss scale, but gradient clipping at 1.0 is not, so the two differ in practice.

• Adaptive density: absolute value or positive part. The text says the density is proportional to |dL/dt| §3.2. Algorithm 4 and the code clamp the derivative at zero, so negative slopes are dropped rather than flipped alg. 4 noise_schedules.py.

• Adaptive hyperparameters. The paper gives the refit period, buffer size, EMA rate, uniform mix and warmup app. B.2. It does not give the spline, knot count, grid or ridge values. The code uses a 1000-point grid, 20 cubic spline knots and ridge 1e-3 log-linear-adaptive.yaml. The repository's default buffer is batch / 10 times the refit period, but the released scripts override it to 50 times the batch [code: configs/noise/log-linear-adaptive.yaml#L6] [code: scripts/train/sudoku/sfm_truncated_adaptive.sh#L28], which matches the paper's "R times batch size" rule. The repository default EMA is 0.0, and the scripts set 0.9 log-linear-adaptive.yaml.

• Time convention. The paper's schedules put alpha = 0 at pure noise (t = 0) §2.1 tab. 5. The code uses the opposite time direction: t = 1 is pure noise and the LogLinear alpha is 1 - t noise_schedules.py, with the cosine-squared alpha written as sin^2(pi (1 - t) / 2) [code: noise_schedules.py#L42-L45]. The two conventions give the same alpha as a function of time, but the time variable is reversed. The sampler's time grid runs from 1 down to eps samplers.py.

• Time sampling and sigma. The paper samples t uniformly on [0, 1] §3. The code draws t from [1e-3, 1] with antithetic sampling config.yaml trainer_base.py. The timestep embedding uses sigma = -log(alpha) trainer_base.py, which the paper does not define.

• MLP form. The paper's MLP formula has no bias terms §3.3 app. B.3. The code's up and down projections have bias terms sphere_arch.py, and the GELU is the tanh approximation sphere_arch.py.

• Temperature grid. The main text says the OWT frontier uses temperatures from 0.70 to 1.20 §4.3. The appendix says 0.50 to 1.20 in steps of 0.05, 15 values app. C.8. The appendix count matches 15 values.

• Optimiser details not in the paper. The code uses a 2500-step constant warmup, gradient clipping at 1.0 and bf16 mixed precision [code: configs/lr_scheduler/constant_warmup.yaml#L2] config.yaml. The paper states none of these for GSM8K or Sudoku. Apply the code values or report them as unknown.

• Sudoku claim versus table. The text says S-FLM and FLM outperform AR and MDLM at every difficulty §4.1. That holds for the two truncated rows of tab. 1. The plain S-FLM row is below MDLM at all three difficulty levels tab. 1.

Pitfalls

• Which truncation side. Eq. 17 gives 0.121 for vocabulary about 50k and d = 768, and the adaptive script keeps the noisy side of the schedule tab. 4 [code: scripts/train/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.sh#L32]. The paper's ablation and the fixed DiT script keep the clean side, with bounds of 0.869 to 0.879 tab. 7 sfm_sphere_dit_truncated_fixed_no_renorm.sh. The paper's own derivation says the posterior becomes almost one-hot once the signal level exceeds alpha* §3.2 app. C.2, which points to the adaptive script's side. Check both against a short run before you trust either.

• Equation 17 is an approximation. The paper derives it from a high-dimension model with uniform random embeddings and an omega = pi/2 approximation app. C.2. The authors say that learned embeddings are likely more structured, and that alpha* "should be treated only as a principled heuristic" §7. The bound is not a guaranteed result.

• Unnormalised embeddings in training. The reported runs do not re-project the embedding table, and the code default is the same [code: configs/model/small-sphere-arch.yaml#L24]. The paper says that leaving the table free anneals its learning rate, because each gradient step grows the norm app. B.5 app. C.4. Re-projecting it after every step, as the code allows, changes training results tab. 6.

• Antipodal and near-collapsed points. The log map is undefined at antipodal points §2.2 (footnote). The code clamps the cosine to 1 - 1e-6 and the sine to at least 1e-6 utils.py utils.py. A latent that sits almost exactly on an embedding gets a near-zero velocity.

• Cost of exact velocity. Eq. 15 sums over the whole vocabulary at each position and each step. Each log map needs a dot product, an arccos and a scaled difference in dimension d. This is O(L · |V| · d) per step, and top-k reduces the sum to k terms §3.1. On TinyGSM the S-arch trains at 2.6 steps per second, against 4.2 for the DiT with the same schedule tab. 3.

• Schedule range also applies at sampling. The sampler reads alpha through the same wrapped noise schedule as training samplers.py, and it integrates from t = 1 to t = eps samplers.py. A truncated alpha range therefore changes the step sizes. Check that a sampling checkpoint is loaded with the schedule it was trained with.

• Single training run per configuration. The authors train one model per configuration and say they could not afford more §7. The reported gaps between rows of tab. 2 are single-seed. Error bars cover the test set only app. C.5.

• Low-temperature gap. At T = 0.1, S-FLM is well behind MDLM and Duo §4.2 app. C.7. The paper gives two likely causes: continuous flows fit a continuous density to discrete data, and the Python solutions have a strongly left-to-right structure §7.

• Sampling steps set the cost. GSM8K results use 1024 steps for diffusion models tab. 2. The released S-FLM scripts default to 32 [code: scripts/sample/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.sh#L12]. Results at 32 steps will not match the table.

• Top-1 is temperature-free. Top-1 velocity does not depend on temperature app. C.8, so it appears as one point on the frontier plot rather than a curve app. C.8 fig. 4.

Lineage

• Builds on: Riemannian flow matching for the sphere §2.3; nGPT for the S-arch backbone §3.3; the cross-entropy training used by CDCD for embedding-space diffusion §3; the discrete diffusion training recipe of MDLM and Duo for architecture, data and evaluation §4; and InfoNoise for the adaptive noise schedule idea §3.2 app. B.2.
• Compared against: AR, MDLM, Duo, CANDI and FLM on Sudoku and GSM8K §4.1 §4.2 tab. 1 tab. 2, and MDLM, Duo and FLM on OpenWebText §4.3 fig. 4.
• Superseded by: not known from content in this wiki.