By Justin Deschenaux and Caglar Gulcehre.
Abstract: Continuous Flow Language Models (FLMs) transport noise to data with a deterministic ODE, avoiding the factorized-sampling assumption of discrete diffusion. But standard FLMs operate on one-hot vectors whose dimension scales with the vocabulary, making them expensive to train. 𝕊-FLM instead operate on the hypersphere, where we transport random points towards the clean embeddings via rotations. Previous FLMs match AR in Generative Perplexity, but high-likelihood samples are not necessarily correct in verifiable domains like math and code; 𝕊-FLM substantially improves over previous FLMs on GSM8K!
This repository contains training and evaluation code for 𝕊-FLM along with the discrete-diffusion / flow-matching baselines we compare against (AR, MDLM, Duo, FLM, CANDI). We release pretrained checkpoints for two settings:
- TinyGSM: math reasoning, SmolLM-135M tokenizer, 250k steps.
- OpenWebText (OWT): general LM, GPT-2 tokenizer, 1M steps.
In addition, we provide training and evaluation scripts (no released checkpoints) for Sudoku (synthetic puzzle dataset, 9-digit vocab + clue-tokens, 20k steps; difficulty levels: easy / medium / hard).
Getting started · Checkpoints · Training · Sampling & evaluation · Citation
Create a fresh environment and install the Python dependencies:
conda create -n sfm python=3.12
conda activate sfm
pip install -r requirements.txtrequirements.txt intentionally does not pin torch or numpy. We work inside the NGC PyTorch container (nvcr.io/nvidia/pytorch:25.02-py3) which already ships matching CUDA / cuDNN / NCCL builds; pip-installing torch on top of that will mess up the NGC build. If you are not using the container, install torch and numpy (we use torch==2.7.0, numpy==1.26.4) before running pip install -r requirements.txt.
All the checkpoints are in the HuggingFace repo jdeschena/s-flm.
tinygsm/ # trained on TinyGSM, 250k steps
ar.ckpt
mdlm.ckpt
duo.ckpt
candi/{lr3e-4,lr1e-3}.ckpt # CANDI uses a smaller learning rate in certain configs, hence we try with both
flm/{default,caps}.ckpt # FLMs use attention softcapping and a custom logits processing. We experiment with and without.
sfm/sphere_dit_truncated_fixed_no_renorm.ckpt # Standard DiT, truncated schedule
sfm/sphere_dit_truncated_adaptive_no_renorm.ckpt # Standard DiT, truncated+adaptive schedule
sfm/sphere_arch_truncated_adaptive_no_renorm.ckpt # S-arch (nGPT-inspired), truncated+adaptive schedule
owt/ # trained on OpenWebText, 1M steps
ar.ckpt
mdlm.ckpt
duo.ckpt
flm.ckpt # Original FLM checkpoint of https://github.com/david3684/flm
sfm.ckpt # S-arch, truncated+adaptive schedule
A single checkpoint:
huggingface-cli download jdeschena/s-flm tinygsm/duo.ckpt \
--local-dir ./checkpointsThe whole repo (or a subset):
# Everything (~47 GB):
huggingface-cli download jdeschena/s-flm --local-dir ./checkpoints
# Just TinyGSM:
huggingface-cli download jdeschena/s-flm --local-dir ./checkpoints \
--include 'tinygsm/**'All training scripts are in scripts/train/. Each script is self-contained and exposes a few environment variables for overrides: OUTPUT_DIR, CACHE_DIR, NUM_NODES, DEVICES. They were originally run on 2 nodes with 4 GPUs (TinyGSM) or 4 nodes with 4 GPUs (OWT).
Sudoku scripts accept a DIFFICULTY environment variable (easy / medium / hard).
Example:
DEVICES=8 NUM_NODES=1 \
OUTPUT_DIR=./outputs/sfm_tinygsm \
bash scripts/train/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.shSampling/eval scripts are in scripts/sample/ and expose the same overrides plus CKPT_PATH and STEPS (default 32).
For TinyGSM, the sampling script uses the GSM8K test and writes a per-example JSON to OUTPUT_DIR; On OWT, we generate unconditional samples and computes generative perplexity.
Example: evaluate an SFM checkpoint on TinyGSM after 64 sampling steps:
CKPT_PATH=./checkpoints/tinygsm/sfm/sphere_arch_truncated_adaptive_no_renorm.ckpt \
STEPS=64 \
OUTPUT_DIR=./eval_runs/sfm_tinygsm \
bash scripts/sample/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.shThe SFM sampling scripts (all 3 TinyGSM SFM scripts and scripts/sample/owt/sfm.sh) expose two extra environment variables:
VELOCITY:exact(default, deterministic velocity) orsample(sample from the velocity distribution).TOPK_VELOCITY:-1(default, full vocab — no top-k filtering) or a positive integer (e.g.1,10) to restrict the velocity to the top-k tokens.
# Default: exact velocity, full vocab (no top-k).
bash scripts/sample/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.sh
# Top-1 velocity, sampled (paper's headline setup).
TOPK_VELOCITY=1 VELOCITY=sample bash scripts/sample/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.sh
# Top-10 velocity, exact.
TOPK_VELOCITY=10 bash scripts/sample/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.sh
# No top-k (full vocab), sampled velocity.
VELOCITY=sample bash scripts/sample/tinygsm/sfm_sphere_arch_truncated_adaptive_no_renorm.shThis codebase builds on a number of excellent open-source projects:
- Duo: discrete diffusion baseline; our overall training/eval scaffolding is descended from theirs.
- FLM: original Flow Language Model implementation; our OWT FLM checkpoint is the released checkpoint from this repo, and our FLM training/sampling code follows it.
- CANDI: continuous-and-discrete diffusion baseline (imported).
- PUMA: reference for the TinyGSM data preparation.
- PRISM: reference for the Sudoku data preparation.
This codebase is distributed under apache-2 license.
@misc{deschenaux2026languagemodelinghypersphericalflows,
title={Language Modeling with Hyperspherical Flows},
author={Justin Deschenaux and Caglar Gulcehre},
year={2026},
eprint={2605.11125},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2605.11125},
}