--- title: "Most LLMs predict the next token. Joey does not." url: https://daily.dev/posts/most-llms-predict-the-next-token-joey-does-not--g43khsvkw source_url: https://daily.dev/posts/most-llms-predict-the-next-token-joey-does-not--g43khsvkw type: freeform source: "Build With GenAI" author: "Lay Sheth aka CLoaKY" published: 2026-06-08T16:17:58.198Z updated: 2026-06-08T16:18:29.493Z tags: ["llm", "pytorch"] reading_time: 4 upvotes: 1 comments: 0 language: en --- > ## Documentation Index > Fetch the complete documentation index at: https://daily.dev/llms.txt > Use this file to discover all available pages before exploring further. # Most LLMs predict the next token. Joey does not. **[Build With GenAI](https://daily.dev/sources/buildwithgenai)** · [@cloaky233](https://daily.dev/cloaky233) · 4 min read · 1 upvotes · 0 comments ## Summary Joey is a 170M-parameter masked diffusion language model (MDLM/LLaDA family) built entirely from scratch — including a custom 16K ByteLevel BPE tokenizer, bidirectional timestep-conditioned Transformer, diffusion loss, and iterative remasking sampler. Unlike autoregressive GPT-style models, Joey starts from a fully masked sequence and iteratively predicts all tokens in parallel, keeping only the most confident predictions and remasking the rest. Trained on ~2B tokens from FineWeb-Edu on an A100 in ~6 hours, then fine-tuned on DailyDialog for conversational SFT. The model produces grammatical, fluent text but lacks global coherence — a known capacity ceiling at this scale. Key lessons came from debugging CUDA OOM errors and repetition collapse, the latter solved by the remasking strategy. Code and weights are publicly available; next steps include scaling to 400M–1B parameters and classifier-free guidance. ## Content ![Generated Image June 07, 2026 - 11_43PM.jpeg](https://media.daily.dev/image/upload/s--PxmZd8GG--/f_auto/v1780935401/ugc/content_0cb13c2d-979a-4e3e-95f2-fee5f08e0625?_a=BAMAMiWQ0) ## Most LLMs predict the next token. Joey does not. GPT-style models are autoregressive: they generate left to right, one token at a time, each token conditioned on the ones before it. Joey belongs to a different family, masked diffusion (the MDLM / LLaDA line of work). Instead of writing left to right, it: 1. Starts from a sequence that is 100% `[MASK]` 2. Predicts every token in parallel 3. Keeps only the tokens it is most confident about 4. Re-masks the rest 5. Repeats until the whole sequence resolves That remasking loop (MaskGIT / LLaDA style) is also what kills the repetition collapse that naive single-pass samplers fall into. In one diagram: ``` FineWeb-Edu -> BPE -> packed token blocks | mask each token w.p. t (forward process, fixed) | bidirectional Transformer(+ t) (reverse process, learned) | 1/t-weighted cross-entropy on masked positions | (after training) all-MASK -> predict, keep confident, re-mask rest -> text (sampling) ``` **Forward process:** corrupt text by replacing tokens with `[MASK]` at a random rate `t`. **Reverse process:** a bidirectional, timestep-conditioned Transformer predicts the originals. **Loss:** cross-entropy on the masked positions only, `1/t`-weighted (the MDLM objective). ## The architecture - Parameters: ~170M - Backbone: bidirectional Transformer (no causal mask), timestep-conditioned - `d_model` / layers / heads: 1024 / 12 / 16 - Context length: 256 tokens - Vocabulary: 16,384 (custom ByteLevel BPE plus `[PAD]` `[BOS]` `[EOS]` `[MASK]`) - MLP: 4x GELU, pre-norm, weight-tied head - Diffusion: masked / absorbing-state (MDLM / LLaDA family) Everything is hand-written: the 16K ByteLevel BPE tokenizer, the bidirectional timestep-conditioned Transformer, the diffusion loss, and the iterative-remasking sampler. No `Trainer`, no pretrained weights. Built test-first, with unit tests for every module. ## Training - Data: FineWeb-Edu, ~2B tokens, own 16K BPE tokenizer - Base: A100-40GB, bf16 + EMA, cosine LR + warmup, 174K steps (~6h), gradient accumulation, hours kill-switch - Fine-tune: DailyDialog, response-only masking (LLaDA-style SFT) - Sampler: remasking (MaskGIT / LLaDA) + repetition penalty + top-p ## The honest result After base training plus conversational fine-tuning, here is an actual exchange: ``` you> Hi! joey> Oh, I am right! It's my favorite, we have always been there for a long time... you> Do you like music? joey> I don't know that much. But I think there is no one... ``` Joey greets correctly, forms grammatical sentences, and holds a conversational register. It is fluent but not yet truly coherent: correct local grammar without sustained global meaning. That is not a bug I gave up on. It is the signature of a capacity ceiling. At 170M parameters the model had essentially converged for its size. It learned _how_ language sounds before it had the room to learn _what_ to actually say. Getting to genuine coherence is primarily a scale problem (more parameters and tokens), and that is the next milestone. ## What actually broke, and what it taught me The two failures I learned the most from: - **CUDA OOM** during training, which forced me to actually understand memory layout, gradient accumulation, and batch packing instead of copying a config. - **Repetition collapse** in sampling, which is where the remasking strategy earns its keep. Naive single-shot decoding loops on itself. Predicting all tokens, keeping only the confident ones, and re-masking the rest breaks the loop. You do not really understand diffusion LLMs until you have debugged your own OOM at 2am and watched a loss curve flatten in front of you. No paper or course gets you there. Building the broken version did. ## Roadmap Done: - From-scratch tokenizer, model, diffusion loss, sampler, training loop - Base pretraining on ~2B tokens plus conversational SFT - Remasking sampler to eliminate repetition loops Next: - Scale up (~400M to 1B) for real coherence (in progress) - Larger, cleaner instruction-tuning data - Classifier-free guidance for conditional sampling - Longer context ## Code and weights - Repo: https://github.com/cloaky233/joey - Weights on Hugging Face: https://huggingface.co/cloaky/joey Built on the shoulders of MDLM (Sahoo et al., 2024), LLaDA (Nie et al., 2025), D3PM (Austin et al., 2021), SEDD (Lou et al., 2024), and MaskGIT (Chang et al., 2022). .... If you have worked with discrete diffusion for text, I would love to hear how you think about the autoregressive vs diffusion tradeoff, especially whether the parallel-decoding speed wins actually survive at scale. ## Similar posts on daily.dev - [Multi-token prediction technique triples LLM inference speed without auxiliary draft models](https://daily.dev/posts/multi-token-prediction-technique-triples-llm-inference-speed-without-auxiliary-draft-models-hxaajvily) · InfoWorld · 1 upvotes · 0 comments --- Tags: [#llm](https://daily.dev/tags/llm), [#pytorch](https://daily.dev/tags/pytorch) [View this post on daily.dev](https://daily.dev/posts/most-llms-predict-the-next-token-joey-does-not--g43khsvkw)