A research-oriented implementation of a diffusion-inspired non-autoregressive language model built entirely from scratch using PyTorch.
This project implements a Transformer-based language model that replaces traditional autoregressive token generation with an iterative mask-and-denoise process inspired by diffusion modeling principles.
Instead of predicting the next token sequentially, the model learns to progressively reconstruct corrupted (masked) text over multiple timesteps.
The goal of this project is to explore an alternative formulation of sequence modeling using diffusion-style iterative refinement.
- Built a Transformer-based architecture from scratch
- Implemented a byte-level BPE tokenizer for text processing
- Designed a timestep-conditioned masking (forward corruption) process
- Developed an iterative denoising-based generation pipeline
- Full training loop using PyTorch + Hugging Face Accelerate
- Modular design for dataset, model, training, and inference components
A custom byte-level BPE tokenizer is used to convert raw text into discrete token sequences.
Input tokens are progressively corrupted using a masking schedule conditioned on diffusion timesteps:
- At higher timesteps → more tokens are masked
- At lower timesteps → less corruption
This simulates a discrete form of noise injection.
A standard Transformer decoder-style architecture is used with:
- Token embeddings
- Positional embeddings
- Timestep embeddings
- Multi-layer self-attention blocks
The model predicts original tokens from partially masked sequences.
Generation is performed via an iterative refinement loop:
- Start from fully or partially masked sequence
- Predict token distributions
- Replace masked tokens
- Repeat for multiple timesteps
- Framework: PyTorch
- Acceleration: Hugging Face Accelerate
- Optimizer: AdamW
- Dataset: TinyStories (subset)
- Objective: Reconstruction of original token sequence from corrupted input
The model is trained on a subset of the TinyStories dataset, which consists of short, simple narrative stories suitable for evaluating small-scale language modeling systems.
This project focuses on:
- Exploring non-autoregressive sequence modeling
- Studying diffusion-style iterative refinement for text generation
- Understanding alternatives to traditional causal language modeling
- The model is able to generate coherent short-form stories after iterative denoising
- Demonstrates stable learning behavior under mask-based corruption training
- Validates feasibility of diffusion-style modeling for discrete text generation
- Does not implement full probabilistic diffusion formulation (e.g., DDPM likelihood framework)
- Limited scale due to compute constraints (TinyStories subset)
- No direct comparison with autoregressive baselines
- Extend to full discrete diffusion formulation
- Add DDIM-style faster sampling
- Compare against GPT-style autoregressive models
- Scale to larger datasets and vocabulary sizes
- Improve generation quality with advanced scheduling strategies
- Python
- PyTorch
- Hugging Face Accelerate
- Tokenizers (BPE)
This project explores a diffusion-inspired alternative to autoregressive language modeling, demonstrating how iterative denoising mechanisms can be applied to discrete token sequences using a Transformer backbone.
This is a research/experimental implementation intended for educational exploration of diffusion-style sequence modeling and is not a production-ready generative language system.