README.md
September 22, 2026 · View on GitHub
Posterior Refinement:
Fast Language Generation via Any-Order Flow Maps
Manan Agarwal*1, Sheel Shah*1, Chanhyuk Lee2, Jaehoon Yoo2, Jerry Huang1,
Seunghoon Hong2, Aditi Raghunathan1, Jinwoo Kim†2, Nicholas M. Boffi†1
1Carnegie Mellon University 2KAIST *Equal contribution †Equal advising
News
- [2026-09] Our paper is accepted in COLM 2026 NonAR-LM Workshop as a Spotlight presentation
Official Code Repository
Posterior Refinement with FMLM+. Posterior Refinement lets the model judge the fit of each token after the fact and fix its own mistakes in parallel. The model generates all tokens in parallel and scores each token's posterior confidence given the entire draft. It commits the high-confidence tokens, re-noises the rest, and repeats.
Abstract
Non-autoregressive generation promises iterative refinement — recursively critiquing, erasing, and regenerating arbitrary subsets of tokens — but existing models fail to realize it. Masked Diffusion Models (MDMs) suffer from factorization error, so sample quality collapses when many tokens are generated simultaneously. Flow Map Language Models (FMLMs) sidestep this bottleneck through joint sequence transport, achieving excellent few-step generation, but sacrifice the inference-time flexibility of MDMs.
We introduce FMLM+, a framework that bridges this gap by equipping FMLM with masking-style noise schedules. While generating the full sequence in a single step, FMLM+ simultaneously scores the global consistency of each token a posteriori. We leverage this to introduce Posterior Refinement, an inference-time strategy that lets the model adaptively self-correct its outputs — matching the performance of discrete baselines with up to 32× fewer NFEs. Across diverse benchmarks, FMLM+ with Posterior Refinement improves the speed–quality tradeoff over both the MDM and FMLM families, providing a scalable foundation for high-fidelity language modeling.
Please refer to our paper for more details on method and results.
How to Run
Install Dependencies
pip install torch>=2.3.0
pip install -r requirements.txt
# Install flash-attn separately matching your python / torch version (see https://github.com/Dao-AILab/flash-attention/releases)
pip install flash-attn==2.8.3 --no-build-isolation
Our DiT backbone supports torch.compile with max-autotune for faster training. Enable it via:
export DIT_USE_COMPILE=TRUE
Training
Update data.cache_dir in each script to point to your dataset location. If the directory is empty, the dataset is automatically generated/downloaded and preprocessed.
| Dataset | Script |
|---|---|
| TinyStories | scripts/train/tinystories.sh |
| OpenWebText | scripts/train/owt.sh |
| TinyGSM | scripts/train/tinygsm.sh |
| Sudoku | scripts/train/sudoku.sh |
| Sudoku (uncond.) | scripts/train/uncond_sudoku.sh |
TinyGSM Variants
For TinyGSM, we provide three training scripts covering the initialization and distillation ablations. The _init and _distill variants load MDLM weights via MDLM_CKPT — set it to your MDLM checkpoint before running.
| Variant | Script | Description |
|---|---|---|
| From scratch | scripts/train/tinygsm.sh | FMLM+ trained from scratch |
| MDLM init | scripts/train/tinygsm_init.sh | FMLM+ initialized from MDLM weights |
| MDLM distillation | scripts/train/tinygsm_distill.sh | FMLM+ distilled from an MDLM teacher (p_distill=0.25) |
Evaluation
Set eval.checkpoint_path in each script to your trained checkpoint before running. Eval uses Posterior Refinement (sampling.schedule=refinement).
| Dataset | Script |
|---|---|
| TinyStories | scripts/eval/tinystories.sh |
| OpenWebText | scripts/eval/owt.sh |
| TinyGSM | scripts/eval/tinygsm.sh |
| Sudoku | scripts/eval/sudoku.sh |
| Sudoku (uncond.) | scripts/eval/uncond_sudoku.sh |
Checkpoints
Pretrained FMLM+ checkpoints are available on Google Drive.
| Dataset | Checkpoint |
|---|---|
| TinyStories | tinystories.ckpt |
| OpenWebText | owt.ckpt |
| TinyGSM | tinygsm.ckpt |
| Sudoku (easy) | sudoku_easy.ckpt |
| Sudoku (medium) | sudoku_medium.ckpt |
| Sudoku (hard) | sudoku_hard.ckpt |
Set eval.checkpoint_path to the downloaded checkpoint path when running an evaluation script (for Sudoku, also set data.difficulty to match the checkpoint).
TinyGSM Variants
The TinyGSM experiments include additional checkpoints for the initialization and distillation ablations:
| Checkpoint | Description |
|---|---|
tinygsm.ckpt | FMLM+ trained from scratch |
tinygsm_init.ckpt | FMLM+ initialized from MDLM weights |
tinygsm_distill.ckpt | FMLM+ distilled from an MDLM teacher (p_distill=0.25) |
tinygsm_mdlm.ckpt | MDLM baseline |
BibTeX
@article{agarwal2026posteriorrefinement,
title={Posterior Refinement: Fast Language Generation via Any-Order Flow Maps},
author={Manan Agarwal and Sheel Shah and Chanhyuk Lee
and Jaehoon Yoo and Jerry Huang and Seunghoon Hong
and Aditi Raghunathan and Jinwoo Kim and Nicholas M. Boffi},
journal={arXiv preprint arXiv:2606.24773},
year={2026},
}
Contact
If you have any questions about the paper, code, or potential collaborations, please feel free to reach out to us at {mananaga, sheels}@cs.cmu.edu.
Acknowledgements
This codebase builds upon FMLM, Duo, and ReDi. We would like to thank Modal Labs for their generous compute grants, which proved invaluable in supporting this work.