ugness/elf-pytorch

This repository is an unofficial PyTorch reproduction of a research paper called ELF, short for Embedded Language Flows, which describes a way of generating text using diffusion models. Diffusion models are a class of AI systems that learn to gradually remove noise from data, a technique first developed for generating images and later adapted for text. The original paper's code was written in JAX, a machine learning framework built by Google, and ran on specialized TPU hardware. This project reimplements the same approach in PyTorch, a more widely used framework, so it can run on standard GPU hardware instead. The key measured result is a generation perplexity of 25.61, a score that reflects how natural the generated text sounds where lower is better, close to the original paper's reported 24.1, achieved after training for about three hours per epoch across eight GPUs. The specific model trained here, called ELF B, has 105 million parameters and generates text without any input prompt, trained on OpenWebText, a large collection of text scraped from the web. The repository includes the model code, training and evaluation scripts built on PyTorch Lightning, which simplifies running training across multiple GPUs, the actual reproduction results and generated text samples per training epoch, and instructions for downloading a pretrained checkpoint so it can be evaluated directly without retraining from scratch. The README notes an important caveat: results are not directly comparable to other baseline models because the data was tokenized and preprocessed differently. This project is aimed at researchers studying or building on the ELF paper and requires familiarity with PyTorch and multi GPU training.