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

TL;DR

Zonkey tackles tokenization bottlenecks and fixed-length generation by introducing a fully differentiable hierarchical diffusion model that progresses from raw characters to document-level representations. It combines Probabilistic Attention with a learnable Segment Splitter, a multi-vector Compressor, a Denoising Diffusion Mixed Model (DDMM), and a differentiable Stitcher to enable end-to-end gradient optimization and emergent word- and sentence-level hierarchies, all operating in a latent space with variable-length outputs. Emergent structure arises without explicit supervision, and diffusion-based generation supports flexible text lengths and non-sequential infilling, enabling scalable long-form synthesis. Trained on Wikipedia with a single GPU, Zonkey demonstrates coherent sentence-level generation and offers potential benefits for domain adaptation and scalable generation; the authors also release source code for reproducibility.

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.

Figures (3)

• Figure 1: Heatmap of the scale matrix adjustment $\log(p_k / p_q)$ in Probabilistic Attention for the example phrase "Alice and Bob". Rows correspond to query positions $q$ (top to bottom), columns to key positions $k$. A value of 1 (red) indicates that the attention score is not affected, and this is the case for all $p_q \geq p_k$ positions. Low values (blue) indicate a very small effect of position $k$ over the post-attention state at position $q$. This antisymmetric modulation softly down-weights contributions from low-existence positions, preserving gradients while approximating hard masking.
• Figure 2: Illustration of existence probabilities and normalized existence shares in the Segment Splitter for overlapping segments from the example phrase ”Alice and Bob”. Left panels: Probabilities for Segment1(full phrase, decaying sharply after ”Alice ”) and Segment 2 (starting at ”and Bob”). These are derived from cumulative products of (1 - p BOS), simulating infinite sequences with natural truncation at low values. Right panels: Corresponding normalized shares for loss calculation, ensuring each position contributes uniformly (summing to 1 across overlaps) despite redundancy in positions 6-12. This mechanism preserves differentiability and incentivizes semantically meaningful splits.
• Figure 3: The DDMM mixed-step training objective. $\mathbf{p}_{1}$: clean compression. $\mathbf{p}_{2}$: heavily noised. $\mathbf{p}_{3}$: denoising and compressing $\mathbf{p}_{2}$, then add light noise. $\mathbf{p}_{4}$: denoising and compressing $\mathbf{p}_{3}$. Loss is determined by the cosine similarity from $\mathbf{p}_{4}$ to its projection on the blue segment $[\mathbf{p}_{1},\mathbf{p}_{2}]$], even though this image uses euclidean distances to ease visualization.