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

arXiv: https://arxiv.org/abs/2601.21768 · HTML: https://mlsanity.com/paper/2601.21768/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Discrete and text diffusion, one of 14 directions in Diffusion models.
- Builds on 24 of its 34 references that are in the map.
- First posted on arXiv 29 Jan 2026.

## Placed in

- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [SimCSE: Simple Contrastive Learning of Sentence Embeddings](https://mlsanity.com/paper/2104.08821/) (2021)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [ByT5: Towards a token-free future with pre-trained byte-to-byte models](https://mlsanity.com/paper/2105.13626/) (2021)
- [DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models](https://mlsanity.com/paper/2210.08933/) (2022)
- [Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling](https://mlsanity.com/paper/2104.06967/) (2021)
- [CANINE: Pre-training an Efficient Tokenization-Free Encoder for Language Representation](https://mlsanity.com/paper/2103.06874/) (2021)
- [Byte Latent Transformer: Patches Scale Better Than Tokens](https://mlsanity.com/paper/2412.09871/) (2024)
- [Hierarchical Transformers for Long Document Classification](https://mlsanity.com/paper/1910.10781/) (2019)
- [Energy-Based Diffusion Language Models for Text Generation](https://mlsanity.com/paper/2410.21357/) (2024)
- [Diffusion Models for Non-autoregressive Text Generation: A Survey](https://mlsanity.com/paper/2303.06574/) (2023)
- [Hierarchical Context Merging: Better Long Context Understanding for Pre-trained LLMs](https://mlsanity.com/paper/2404.10308/) (2024)
- [TESS 2: A Large-Scale Generalist Diffusion Language Model](https://mlsanity.com/paper/2502.13917/) (2025)
- [Improving Transformers with Probabilistic Attention Keys](https://mlsanity.com/paper/2110.08678/) (2021)
- [Probabilistic Attention for Interactive Segmentation](https://mlsanity.com/paper/2106.15338/) (2021)
- [Rethinking Tokenization: Crafting Better Tokenizers for Large Language Models](https://mlsanity.com/paper/2403.00417/) (2024)
- [Next Semantic Scale Prediction via Hierarchical Diffusion Language Models](https://mlsanity.com/paper/2510.08632/) (2025)

## Abstract (arXiv)

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.