Update README.md

README.md

@@ -12,4 +12,194 @@ tags:
 - gpt2
 - research
 pipeline_tag: text-generation
----
+---
+
+# Diffusion-LLM
+
+> **A small, experimental discrete diffusion language model built on GPT-2.**
+
+This repository contains checkpoints from an experiment in iterative text
+denoising. Instead of predicting only the next token from a left-to-right
+context, the model receives a corrupted sequence and predicts the original
+tokens across the whole sequence with bidirectional attention.
+
+The checkpoints are intended for research and reproducibility. They are not
+drop-in `transformers` causal language models and should not be treated as
+production-ready text generators.
+
+## What is inside?
+
+The model starts from [`openai-community/gpt2`](https://huggingface.co/openai-community/gpt2)
+and adapts its Transformer backbone for discrete diffusion-style denoising:
+
+- GPT-2 embeddings and Transformer blocks are retained.
+- Causal attention is replaced with bidirectional attention.
+- A learned diffusion-timestep embedding is added to every token embedding.
+- The language-model head predicts the clean token at each position.
+- Denoising is performed by repeatedly updating corrupted positions.
+
+The training code adds two special tokens to the GPT-2 tokenizer:
+`<|pad|>` and `<|mask|>`. The vocabulary is resized before training, so the
+same tokenizer construction is required when loading a checkpoint.
+
+## Checkpoints
+
+The uploaded weights follow this layout:
+
+| Path | Training corruption | Refinement steps during training |
+| --- | --- | ---: |
+| `default.pt` | Initial/reference checkpoint | — |
+| `1_step/similar_100.pt` | Semantically similar replacements | 1 |
+| `1_step/mask_100.pt` | Mask-token corruption | 1 |
+| `1_step/random_100.pt` | Random-token replacement | 1 |
+| `1_step/mix_30_30_40.pt` | 30 training epochs similar → 30 mask → 40 random | 1 |
+| `4_step/similar_100_4_steps.pt` | Semantically similar replacements | 4 |
+| `4_step/mask_100_4_steps.pt` | Mask-token corruption | 4 |
+| `4_step/random_100_4_steps.pt` | Random-token replacement | 4 |
+| `4_step/mix_30_30_40_4_steps.pt` | 30 training epochs similar → 30 mask → 40 random | 4 |
+
+All listed experimental runs use the same main setup: GPT-2, 100 diffusion
+timesteps, maximum sequence length 64, learning rate `1e-5`, and WikiText-2
+(`Salesforce/wikitext`, `wikitext-2-raw-v1`). The 4-step runs use a rollout
+loss decay of `0.5`.
+
+## Corruption strategies
+
+The training pipeline supports three discrete corruption operators:
+
+1. **Similar** — replaces tokens with nearby tokens in the GPT-2 embedding
+   space. This creates difficult, semantically close mistakes.
+2. **Mask** — replaces tokens with `<|mask|>`, the classic masked-denoising
+   condition.
+3. **Random** — replaces tokens with random vocabulary entries.
+
+The mixed checkpoints use these methods as a sequential training curriculum:
+**30 epochs of similar-token corruption, 30 epochs of masking, and 40 epochs
+of random replacement**. Corruption probability is sampled between `0.01` and
+`0.95`; the similar-token operator uses 20 nearest neighbours.
+
+## Quick start
+
+Clone this repository so that the model wrapper and corruption utilities are
+available alongside the downloaded checkpoint:
+
+```bash
+git clone https://huggingface.co/AndreiRabau/diffusion-llm
+cd diffusion-llm
+pip install torch transformers
+```
+
+The following example loads the 4-step mixed checkpoint and denoises a
+sequence containing mask tokens:
+
+```python
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer
+
+from model_wrappers import GPT2DiffusionTransformer
+
+
+MODEL_NAME = "openai-community/gpt2"
+CHECKPOINT = Path("4_step/mix_30_30_40_4_steps.pt")
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+tokenizer.add_special_tokens({
+    "pad_token": "<|pad|>",
+    "mask_token": "<|mask|>",
+})
+
+model = GPT2DiffusionTransformer.from_file_path(
+    file_path=CHECKPOINT,
+    model_name=MODEL_NAME,
+    num_diffusion_steps=100,
+    vocabulary_size=len(tokenizer),
+    device="cpu",
+).eval()
+
+text = "Diffusion models can refine [MASK] sequences over several steps."
+inputs = tokenizer(text, return_tensors="pt")
+
+# For a real masked input, create the sequence with tokenizer.mask_token_id.
+# Here we replace one token to keep the example self-contained.
+corrupted_ids = inputs["input_ids"].clone()
+corrupted_positions = torch.zeros_like(corrupted_ids, dtype=torch.bool)
+corrupted_ids[0, 4] = tokenizer.mask_token_id
+corrupted_positions[0, 4] = True
+
+reconstructed_ids = model.denoise(
+    corrupted_ids=corrupted_ids,
+    attention_mask=inputs["attention_mask"],
+    corrupted_positions=corrupted_positions,
+    num_iterations=4,
+)
+
+print(tokenizer.decode(reconstructed_ids[0], skip_special_tokens=True))
+```
+
+The `.pt` files contain a PyTorch `state_dict`. Loading them therefore
+requires the repository's `GPT2DiffusionTransformer` implementation and the
+same model configuration used during training (`num_diffusion_steps=100` and
+the resized tokenizer vocabulary).
+
+## Training and evaluation details
+
+Training uses cross-entropy only on corrupted, non-padding positions. For
+multi-step training, each step feeds the model's predictions back into the
+next step, exposing the model to its own errors. At evaluation time, the
+denoiser uses confidence-guided refinement: the most confident corrupted
+positions are updated first, while less certain positions remain available for
+later iterations. This difference is intentional.
+
+The included experiments evaluate on the `test` split of
+[`cimec/lambada`](https://huggingface.co/datasets/cimec/lambada) with
+corruption rates from 25% to 95% and 2, 4, 5, 8, 10, 20, or 50 denoising
+iterations, depending on the run. Evaluation artifacts are stored in
+`results/eval/` in the source project.
+
+## Intended use
+
+This release is useful for:
+
+- studying corruption policies for discrete diffusion language models;
+- comparing one-step and iterative denoising rollouts;
+- experimenting with confidence-based token refinement;
+- reproducing the accompanying small-scale WikiText-2 experiments.
+
+It is **not** intended for factual question answering, safety-critical
+generation, deployment, or direct comparison with large-scale diffusion LMs.
+The checkpoints inherit the limitations of GPT-2 and the small experimental
+training setup, including potential memorization, bias, and unstable outputs.
+
+## Limitations and open questions
+
+- These are research checkpoints rather than a fully packaged inference
+  pipeline.
+- The `.pt` format stores weights only; configuration and tokenizer files are
+  reconstructed by the loader.
+- The model was trained on short sequences (maximum length 64), so longer
+  contexts are outside the validated setup.
+- The corruption curriculum and the semantic-neighbour operator have not been
+  established as generally superior to standard masked or random corruption.
+- Reported results should be interpreted as exploratory measurements, not as
+  a leaderboard claim.
+
+## Citation
+
+If you use these checkpoints or the training setup, please cite this
+repository and the upstream GPT-2 work:
+
+```bibtex
+@misc{diffusion_llm_experiment,
+  title  = {Diffusion-LLM: Experimental GPT-2 Discrete Denoising Checkpoints},
+  author = {Andrei Rabau},
+  year   = {2026},
+  url    = {https://huggingface.co/AndreiRabau/diffusion-llm}
+}
+```
+
+The implementation and experiment configurations are available in the source
+repository, including [`model_wrappers/gpt2_diffusion_transformer_wrapper.py`](model_wrappers/gpt2_diffusion_transformer_wrapper.py),
+[`train_utils/train.py`](train_utils/train.py), and the files under
+[`experiments/`](experiments/).