jdeschena/s-flm

S-FLM is a research codebase for a new kind of AI language model, the type of system that generates text, solves math problems, or completes code. Most language models predict one word at a time, moving left to right. S-FLM explores a different approach: it uses a process called flow matching to gradually transform random noise into finished text. The key idea is that words and sentences can be represented as points on a mathematical sphere. The model learns to rotate random starting points toward the correct text, rather than picking words one at a time. This spherical approach gives the S in S-FLM. The authors state that earlier flow based models matched standard left to right language models at producing plausible sounding text, but fell short on tasks where getting the exact answer right matters, such as math problems. S-FLM shows improvement on the GSM8K math reasoning benchmark compared to those earlier flow based approaches. The repository contains the training and evaluation code from the accompanying research paper, along with pretrained checkpoints that can be downloaded from HuggingFace. It supports training on two main datasets: TinyGSM, a math reasoning dataset, and OpenWebText, a general language modeling dataset. A synthetic Sudoku puzzle task with adjustable difficulty is also included for controlled experiments. Training scripts cover several competing model types for comparison, including standard autoregressive models and other diffusion or flow based approaches named MDLM, Duo, FLM, and CANDI, which makes it possible to reproduce the paper's side by side comparisons. Setup requires a Python 3.12 environment with specific torch and numpy versions, and the authors recommend using an NVIDIA container image to keep CUDA versions aligned. The project is written in Python.