Zonkey: A Hierarchical Diffusion Language Model with Differentiable Tokenization and Probabilistic Attention

TL;DR

Zonkey introduces a hierarchical diffusion language model with a differentiable tokenizer and probabilistic attention, enabling end-to-end training from raw characters to document-level representations, improving adaptability and scalability of language models.

Contribution

It presents a fully trainable pipeline with a differentiable tokenizer and probabilistic attention, allowing end-to-end optimization of language models from raw text inputs.

Findings

Generates coherent, variable-length text from noise.

Emergent hierarchies align with linguistic structures.

Outperforms entropy-based tokenizers in data distribution alignment.

Abstract

Large language models (LLMs) have revolutionized natural language processing, yet they remain constrained by fixed, non-differentiable tokenizers like Byte Pair Encoding (BPE), which hinder end-to-end optimization and adaptability to noisy or domain-specific data. We introduce Zonkey, a hierarchical diffusion model that addresses these limitations through a fully trainable pipeline from raw characters to document-level representations. At its core is a differentiable tokenizer (Segment Splitter) that learns probabilistic beginning-of-sequence (BOS) decisions, enabling adaptive splits that emerge as linguistically meaningful (e.g., word boundaries at spaces, sentence starts at periods) without explicit supervision. This differentiability is enabled by our novel Probabilistic Attention mechanism, which incorporates position-specific existence probabilities to simulate soft masking over theoretically infinite sequences while preserving gradients. Sequences decay probabilistically rather than relying on end-of-sequence tokens, supporting variable-length outputs. Hierarchical levels compress sequences into higher abstractions (e.g., character n-grams to word-like vectors, then sentence-like), with reconstruction via our Denoising Diffusion Mixed Model (DDMM) for stable and efficient denoising in latent space. A Stitcher ensures overlap invariance across segments. Trained end-to-end on Wikipedia, Zonkey generates coherent, variable-length text from noise, demonstrating emergent hierarchies and promising qualitative alignment to data distributions compared to entropy-based learnable tokenizers. Our approach advances toward fully gradient-based LLMs, with potential for better domain adaptation and scalable generation. We release the source code for training and reproducing our experiments.