| --- |
| license: mit |
| language: |
| - en |
| library_name: pytorch |
| tags: |
| - diffusion-language-model |
| - masked-diffusion |
| - mdlm |
| - llada |
| - text-generation |
| - from-scratch |
| pipeline_tag: text-generation |
| --- |
| |
| # Joey 🐤 — a diffusion language model, from scratch |
|
|
| Joey is a ~170M-parameter **masked / absorbing-state diffusion language model** (MDLM / LLaDA |
| family), implemented from scratch in PyTorch. Instead of generating left-to-right like GPT, it |
| generates by **iterative denoising**: starting from a fully `[MASK]`ed sequence and progressively |
| unmasking, re-deciding low-confidence tokens along the way. |
|
|
| **Code & full write-up:** https://github.com/CLoaKY233/joey |
|
|
| > **Status: work in progress.** This checkpoint is a small base + conversational fine-tune. It is |
| > fluent but capacity-limited — it learns grammar and conversational register, not sustained meaning. |
| > Scaling up is the next milestone. |
|
|
| ## Files |
|
|
| | File | Description | |
| |---|---| |
| | `joey_chat.pt` | Conversational model (base + DailyDialog SFT) — use this to chat | |
| | `joey_base.pt` | Base model after pretraining (step 174k) | |
| | `tok.json` | The 16K ByteLevel BPE tokenizer | |
|
|
| ## Model details |
|
|
| | Property | Value | |
| |---|---| |
| | Parameters | ~170M | |
| | Backbone | Bidirectional Transformer (no causal mask), timestep-conditioned | |
| | `d_model` / layers / heads | 1024 / 12 / 16 | |
| | Context length | 256 | |
| | Vocabulary | 16,384 (custom ByteLevel BPE) | |
| | Objective | Masked diffusion, `1/t`-weighted cross-entropy on masked positions | |
| | Training data | FineWeb-Edu (~2B tokens) | |
| | Fine-tuning | DailyDialog, response-only masking (LLaDA-style SFT) | |
| | Sampler | Remasking (MaskGIT/LLaDA) + repetition penalty + top-p | |
|
|
| ## Usage |
|
|
| Clone the [code repo](https://github.com/CLoaKY233/joey), place `joey_chat.pt` and `tok.json` in |
| `artifacts/`, then: |
|
|
| ```bash |
| uv run python scripts/chat.py |
| ``` |
|
|
| ## References |
|
|
| - Sahoo et al., *Simple and Effective Masked Diffusion Language Models* (MDLM), NeurIPS 2024 |
| - Nie et al., *Large Language Diffusion Models* (LLaDA), 2025 |
| - Chang et al., *MaskGIT: Masked Generative Image Transformers*, CVPR 2022 |
|
|
| ## License |
|
|
| MIT |
|
|
