Upload README.md with huggingface_hub

README.md

@@ -1,3 +1,125 @@
 ---
-license: apache-2.0
+license: mit
+tags:
+  - language-model
+  - multi-token-prediction
+  - push-forward-language-model
+  - text-generation
+  - distillation
+datasets:
+  - lm1b
+  - openwebtext
+arxiv: "2606.10820"
 ---
+
+# K-Forcing: Joint Next-K-Token Decoding via Push-Forward Language Modeling
+
+<p align="center">
+  <a href="https://arxiv.org/abs/2606.10820"><img src="https://img.shields.io/badge/arXiv-2606.10820-b31b1b.svg" alt="arXiv"></a>
+  <a href="https://github.com/Tangzw2020/K-Forcing"><img src="https://img.shields.io/badge/GitHub-Code-blue?logo=github" alt="GitHub"></a>
+</p>
+
+## Overview
+
+K-Forcing distills an autoregressive (AR) language model into a **push-forward language model (PFLM)** that generates **k tokens in one forward pass**. It maps k independent uniform noise variables to k future tokens jointly via an inverse-CDF construction, enabling fixed-length multi-token decoding that is fully compatible with standard KV-cache batch serving.
+
+**Key results**: ~2.4–3.5× batch-serving throughput speedup at modest quality degradation on LM1B and OpenWebText with ~100M-param Transformers.
+
+## Checkpoints
+
+This repository contains four checkpoints:
+
+| File | Model | Dataset | Parameters | Description |
+|------|-------|---------|------------|-------------|
+| `ar_openwebtxt.ckpt` | AR | OpenWebText | ~100M | Autoregressive teacher model (GPT-2 tokenizer, seq_len=1024) |
+| `ar_best_lm1b.ckpt` | AR | LM1B | ~100M | Autoregressive teacher model (custom tokenizer, seq_len=128) |
+| `pflm_owt_k4.ckpt` | PFLM (k=4) | OpenWebText | ~100M | Push-forward LM, decodes 4 tokens per forward pass |
+| `pflm_lm1b_k4.ckpt` | PFLM (k=4) | LM1B | ~100M | Push-forward LM, decodes 4 tokens per forward pass |
+
+All models share a 12-layer causal Transformer backbone (768 hidden dim, 12 heads), following the architecture from [MDLM](https://arxiv.org/abs/2406.07524) (Sahoo et al., 2024).
+
+## Download
+
+```python
+from huggingface_hub import hf_hub_download
+
+# Download a specific checkpoint
+ckpt_path = hf_hub_download(
+    repo_id="zwave/K-Forcing",
+    filename="pflm_owt_k4.ckpt",  # or: ar_openwebtxt.ckpt, ar_best_lm1b.ckpt, pflm_lm1b_k4.ckpt
+)
+```
+
+Or download all checkpoints at once:
+```python
+from huggingface_hub import snapshot_download
+
+snapshot_download(repo_id="zwave/K-Forcing", local_dir="./checkpoints")
+```
+
+Or via CLI:
+```bash
+huggingface-cli download zwave/K-Forcing --local-dir ./checkpoints
+```
+
+## Usage
+
+Clone the [K-Forcing repository](https://github.com/Tangzw2020/K-Forcing) and follow setup instructions there:
+
+```bash
+git clone https://github.com/Tangzw2020/K-Forcing.git
+cd K-Forcing
+
+# Setup environment
+mkdir -p wheels
+wget -P wheels https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.6/flash_attn-2.5.6+cu122torch2.2cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
+uv sync
+```
+
+### AR Inference
+
+```bash
+python batch_inference_with_prefix.py \
+    --model ar --task owt \
+    --ckpt_path ./checkpoints/ar_openwebtxt.ckpt \
+    --prefix_file assets/prefix_owt_examples.jsonl \
+    --batch_size 4 --n_per_prefix 1
+```
+
+### PFLM Inference (K=2 tokens per forward pass)
+
+```bash
+python batch_inference_with_prefix.py \
+    --model pflm --task owt \
+    --ckpt_path ./checkpoints/pflm_owt_k4.ckpt \
+    --prefix_file assets/prefix_owt_examples.jsonl \
+    --batch_size 4 --n_per_prefix 1 --K 2 --freq_penalty 0.3
+```
+
+The PFLM checkpoint trained with k=4 supports inference with any K ≤ 4.
+
+## Architecture
+
+- **Backbone**: 12-layer causal Transformer (~100M params), 768 hidden dim, 12 heads
+- **Noise encoder**: sinusoidal + MLP, encodes each Uniform(0,1) noise variable into a token embedding
+- **Fully causal design**: noise tokens attend causally — each zⱼ sees context + z₁..zⱼ
+- **Shared prediction head**: same linear head as AR, applied at each noise-token position
+- **Training**: progressive self-forcing distillation (AR → k=1 → k=2 → k=4)
+
+## Citation
+
+```bibtex
+@misc{tang2026kforcingjointnextktokendecoding,
+      title={K-Forcing: Joint Next-K-Token Decoding via Push-Forward Language Modeling},
+      author={Zhiwei Tang and Yuanyu He and Yizheng Han and Wangbo Zhao and Jiasheng Tang and Fan Wang and Bohan Zhuang},
+      year={2026},
+      eprint={2606.10820},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2606.10820},
+}
+```
+
+## License
+
+This project is licensed under the MIT License.