xLM: A Python Package for Non-Autoregressive Language Models
Paper · 2026 · Dhruvesh Patel, Durga Prasad Maram, Sai Sreenivas Chintha et al.
Summary
Research on non-autoregressive text generation is growing, including masked diffusion, Gaussian diffusion, insertion-based and edit-based language models. Unlike autoregressive LMs, these approaches lack common tooling. Most implementations are one-off, and each method needs its own collation, loss and decoding code, so fair comparisons are hard. The authors present xLM, a Python package built on PyTorch, PyTorch Lightning and Hydra. It aims to speed up prototyping of small non-autoregressive LMs.
Its design follows a "maximal independence" philosophy:
- Shared core components (Harness, DataModule) delegate model-specific work to pluggable parts.
- Each model lives in a self-contained folder that can be copied from a template.
- Hydra configs let users swap whole components without editing code.
A walkthrough shows how to build the Insertion Language Model (ILM) step by step. To check fidelity, the authors re-implement several models on star-graph planning tasks and on LM1B. Results stay close to earlier published figures. A companion package, xlm-models, is planned to provide small pre-trained models for the community.
Key points
- Core architecture: a Harness (a LightningModule) holds the Model, LossFunction, Predictor and Metrics slots. A TextDataModule manages one DatasetManager per dataset, mapped many-to-many to the train/val/test/predict stages.
- Training setups: single-GPU, multi-GPU and multi-node, with both map-style and iterable datasets.
- Workflows: a single
xlmcommand runs train, eval and generate jobs, plus adebug=overfitmode. Models are discovered throughXLM_MODELS_PACKAGESorXLM_MODELS_PATH. - Star-graph planning benchmark (StarEasy and harder variants): ARLM, MLM, MDLM and ILM results fall within 2% of the original papers. On the Hard variant, ILM reaches 97.5% sequence accuracy, while the other models score roughly 21–25%.
- LM1B benchmark: 12-layer ARLM, MDLM and ILM models are scored by negative log-likelihood under Llama 3.2, and the outcomes are close to those reported by Patel et al. (2025). Planned extensions include edit-based models and non-text tasks such as molecule generation.
Institutions
Fields, topics and concepts
📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works
🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works
Cite
Cite as Dhruvesh Patel, Durga Prasad Maram, Sai Sreenivas Chintha, Benjamin Rozonoyer and Andrew McCallum. 2026. xLM: A Python Package for Non-Autoregressive Language Models. Proceedings of the 19th Conference of the European Chapter of the Association for Computational Linguistics (Volume 3: System Demonstrations).
@inproceedings{patel2026xlm,
title = {{xLM: A Python Package for Non-Autoregressive Language Models}},
author = {Dhruvesh Patel, Durga Prasad Maram, Sai Sreenivas Chintha, Benjamin Rozonoyer and Andrew McCallum},
booktitle = {Proceedings of the 19th Conference of the European Chapter of the Association for Computational Linguistics (Volume 3: System Demonstrations)},
year = {2026},
doi = {10.18653/v1/2026.eacl-demo.31},
url = {https://doi.org/10.18653/v1/2026.eacl-demo.31}
}References
- Gregor Bachmann and Vaishnavh Nagarajan. 2024. The pitfalls of next-token prediction. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 2296–2318. PMLR. find
- Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. 2014. One billion word benchmark for measuring progress in statistical language modeling. Preprint, arXiv:1312.3005.
- Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. 2024. Flex attention: A programming model for generating optimized attention kernels. Preprint, arXiv:2412.05496. find
- William Falcon and The PyTorch Lightning team. 2019. PyTorch Lightning. find
- Erich Gamma, Richard Helm, Ralph Johnson, and John Vlissides. 1995. Design patterns: elements of reusable object-oriented software. Addison-Wesley Longman Publishing Co., Inc., USA. find
- Matt Gardner, Joel Grus, Mark Neumann, Oyvind Tafjord, Pradeep Dasigi, Nelson F. Liu, Matthew Peters, Michael Schmitz, and Luke S. Zettlemoyer. 2017. Allennlp: A deep semantic natural language processing platform.
- Ishaan Gulrajani and Tatsunori Hashimoto. 2023. Likelihood-based diffusion language models. In Thirty-seventh Conference on Neural Information Processing Systems. find
- Marton Havasi, Brian Karrer, Itai Gat, and Ricky T. Q. Chen. 2025. Edit Flows: Flow Matching with Edit Operations. Preprint, arXiv:2506.09018. — [arXiv:2506.09018](https://arxiv.org/abs/2506.09018) find
- John J Irwin and Brian K Shoichet. 2005. Zinc- a free database of commercially available compounds for virtual screening. Journal of chemical information and modeling, 45(1):177–182. find
- Jaeyeon Kim, Lee Cheuk-Kit, Carles Domingo-Enrich, Yilun Du, Sham Kakade, Timothy Ngotiaoco, Sitan Chen, and Michael Albergo. 2025. Any-Order Flexible Length Masked Diffusion. Preprint, arXiv:2509.01025. find
- Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. 2023. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles.
- Pengfei Li, Qichang Zheng, and Ziyi Jiang. 2025. An empirical study on the accuracy of large language models in api documentation understanding: A cross-programming language analysis. Journal of Computing Innovations and Applications, 3(2):1–14. find
- Lightning-AI. 2023. Litgpt. https://github.com/ Lightning-AI/litgpt. ↗
- Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bhagia, Yuling Gu, Shengyi Huang, Matt Jordan, Nathan Lambert, Dustin Schwenk, Oyvind Tafjord, Taira Anderson, David Atkinson, Faeze Brahman, Christopher Clark, Pradeep Dasigi, Nouha Dziri, and 21 others. 2024. 2 ol… find
- Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. 2019. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demonstr…
- Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Te-jani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, and 2 others. 2019.… — PyTorch: An Imperative Style, High-Performance Deep Learning Library
- Dhruvesh Patel, Aishwarya Sahoo, Avinash Amballa, Tahira Naseem, Tim GJ Rudner, and Andrew Mc-Callum. 2025. Insertion language models: Sequence generation with arbitrary-position insertions. arXiv preprint arXiv:2505.05755. — [arXiv:2505.05755](https://arxiv.org/abs/2505.05755) find
- William Peebles and Saining Xie. 2023. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4195–4205.
- Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. 2012. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52(11):2864–2875. find
- Subham Sekhar Sahoo, Marianne Arriola, Aaron Gokaslan, Edgar Mariano Marroquin, Alexander M. Rush, Yair Schiff, Justin T. Chiu, and Volodymyr Kuleshov. 2024. Simple and Effective Masked Diffusion Language Models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.
- Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. 2024. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063.
- Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, Shengyi Huang, Kashif Rasul, and Quentin Gallouédec. 2020. Trl: Transformer reinforcement learning. https://github.com/huggingface/trl. ) return batch A.2 ↗
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/patel-2026-xlm-python-package-non-autoregressive/