# ELF: Embedded Language Flows (Unofficial PyTorch Reproduction)

> [!CAUTION]
> **! Caution !**
>
> The results are not directly comparable with baselines ([MDLM](https://github.com/kuleshov-group/mdlm), [Duo](https://github.com/s-sahoo/duo), [FLM](https://github.com/david3684/flm), ...)
> due to tokenization and preprocessing differences used in the ELF paper.
>
> Specifically, ELF uses a custom preprocessed OpenWebText dataset (see [`openwebtext-t5`](https://huggingface.co/datasets/embedded-language-flows/openwebtext-t5)).
> This is tokenized with the T5 tokenizer, not the GPT-2 tokenizer which is used in the standard setting in the literature. In addition, the paper's preprocessing pipeline includes a custom packing scheme with full details not disclosed in the paper.

---

> **This is an unofficial PyTorch reproduction** of *ELF: Embedded Language Flows*.
> It is not affiliated with or endorsed by the paper authors. The official JAX/TPU
> implementation is at <https://github.com/lillian039/ELF>, and the official
> checkpoints are in HuggingFace at
> [`embedded-language-flows`](https://huggingface.co/embedded-language-flows).
>
> This repository was developed using [Claude Code](https://claude.com/claude-code).

## Reproduction status

OpenWebText (unconditional), ELF-B (105M), 32-step SDE, γ=1.5, SC-CFG=3:

| Metric | Paper (TPU v5p-64) | Reproduction (8× B200 DDP, Lightning) |
| --- | --- | --- |
| Gen. PPL ↓ | 24.1 | **25.61** |
| Entropy ↑ | 5.15 | **5.20** |

Per-epoch results (32-step SDE, 256 samples):

| Epoch | Step | Gen. PPL | Entropy |
| --- | --- | --- | --- |
| 1 | 38 034  | 2.73¹  | 0.70¹ |
| 2 | 76 068  | 37.11  | 5.17 |
| 3 | 114 102 | 28.63  | 5.21 |
| 4 | 152 136 | 25.00  | 5.16 |
| 5 | 190 170 | 25.58  | 5.19 |
| 6 | 228 204 | 26.11  | 5.21 |

All samples used for the measurements can be found in
[`reproduction/elf_b-owt/eval1000/metrics.jsonl`](reproduction/elf_b-owt/eval1000/metrics.jsonl)
and [`reproduction/elf_b-owt/per_epoch/metrics.jsonl`](reproduction/elf_b-owt/per_epoch/metrics.jsonl).

## TODO
- [ ] Train ELF and/or some of the baselines ([MDLM](https://github.com/kuleshov-group/mdlm), [Duo](https://github.com/s-sahoo/duo), [FLM](https://github.com/david3684/flm), ...) in a directly comparable setting (https://huggingface.co/datasets/Skylion007/openwebtext).

## What's in this repo

- [`pytorch_lightning/`](pytorch_lightning/): model, training
  script (`train_lightning.py`), eval (`eval_lightning.py`), and
  utilities. 8-GPU CUDA DDP via PyTorch Lightning.
- [`reproduction/elf_b-owt/`](reproduction/elf_b-owt/): config snapshot, 1000 final
 samples, and per-epoch samples. The
  checkpoint is hosted separately (see [Quickstart](#quickstart-evaluate-the-reproduced-checkpoint)).

## Quickstart — evaluate the reproduced checkpoint

```bash
# 1. Environment (conda)
conda env create -f environment.yml -n elf-pytorch && conda activate elf-pytorch

# 2. Download the reproduced final EMA checkpoint (1.4 GB)
pip install huggingface_hub
huggingface-cli download Ugness/elf-torch last.ckpt \
    --local-dir reproduction/elf_b-owt/

# 3. Run the 1000-sample evaluation
cd pytorch_lightning/
torchrun --nproc_per_node=8 --master_port=29510 eval_lightning.py \
    --config configs/training_configs/train_owt_ELF-B.yml \
    --checkpoint_path ../reproduction/elf_b-owt/last.ckpt \
    --num_samples 1000
# Expected: Gen. PPL ≈ 25.6, sample entropy ≈ 5.20.
```

### Per-epoch checkpoints

The checkpoints are under this HF repo:
[`checkpoints/`](https://huggingface.co/Ugness/elf-torch/tree/main/checkpoints).
```bash
# Example: pull epoch 4 ckpt.
huggingface-cli download Ugness/elf-torch \
    checkpoints/checkpoint_epoch03_step00152136.ckpt \
    --local-dir reproduction/elf_b-owt/
```

## Quickstart — train from scratch

```bash
cd pytorch_lightning/
torchrun --nproc_per_node=8 --master_port=29501 train_lightning.py \
    --config configs/training_configs/train_owt_ELF-B.yml
```

## Reproduction details

- **Hardware:** 8× NVIDIA B200 (sm_100), CUDA 12.8.
  `broadcast_buffers=False`. See `pytorch_lightning/train_lightning.py`.
- **Wall-clock:** ~3 hours per epoch.


### Differences vs the paper run

| Aspect | Paper | This reproduction |
| --- | --- | --- |
| Hardware | TPU v5p-64 | 8× B200 DDP |
| Framework | JAX/Flax | PyTorch Lightning |
| Epochs | 5 | 6 (one extra to reach entropy ≈ 5.20) |
| Optimizer / objective | Muon + L2 denoise + CE decode (decoder_prob=0.2) | Unchanged |
| Schedule, noise scale, time schedule, SC, CFG | Unchanged | Unchanged |