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

• URL: http://arxiv.org/abs/2601.21768v1
• Date: Thu, 29 Jan 2026 14:17:37 GMT
• Title: Zonkey: A Hierarchical Diffusion Language Model with Differentiable Tokenization and Probabilistic Attention
• Authors: Alon Rozental,
• Abstract summary: Zonkey is a hierarchical diffusion model that addresses limitations through a fully trainable pipeline from raw characters to document-level representations.<n>At its core is a differentiable tokenizer that learns probabilistic beginning-of-sequence (BOS) decisions.<n>Zonkey generates coherent, variable-length text from noise, demonstrating emergent hierarchies.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Token Maturation: Autoregressive Language Generation via Continuous Token Dynamics [0.7252027234425333]
  We introduce a continuous autoregressive formulation of language generation in which tokens are represented as continuous vectors that emphmature over multiple update steps before being discretized.<n>We show that this maturation process alone is sufficient to produce coherent and diverse text using deterministic decoding (argmax)<n>Additional perturbations, such as dynamics or history smoothing, can be incorporated naturally but are not required for the model to function.
  arXiv  Detail & Related papers  (2026-01-08T11:44:34Z)
• Continuous Autoregressive Language Models [56.49239051750678]
  We introduce Continuous Autoregressive Language Models (CALM)<n>CALM uses a high-fidelity autoencoder to compress a chunk of K tokens into a single continuous vector.<n>We develop a comprehensive likelihood-free framework that enables robust training, evaluation, and controllable sampling.
  arXiv  Detail & Related papers  (2025-10-31T17:58:11Z)
• Semantic Fusion with Fuzzy-Membership Features for Controllable Language Modelling [0.0]
  semantic fusion is a lightweight scheme that augments a Transformer language model (LM) with a fuzzy-membership feature channel.<n>Each token is represented by a vector of interpretable features whose values are graded degrees from differentiable membership functions.<n>This approach adds only small overhead, remains fully compatible with tied input-output embeddings, and provides an interpretable pathway for conditioned natural language generation.
  arXiv  Detail & Related papers  (2025-09-14T22:11:09Z)
• FLEXITOKENS: Flexible Tokenization for Evolving Language Models [9.003053181721823]
  Language models (LMs) are challenging to adapt to new data distributions by simple finetuning.<n>This is due to the rigidity of their subword tokenizers, which typically remain unchanged during adaptation.<n>We develop byte-level LMs with learnable tokenizers to make tokenization adaptive.
  arXiv  Detail & Related papers  (2025-07-17T01:55:41Z)
• Sampling from Your Language Model One Byte at a Time [82.71473348639489]
  Tokenization can introduce distortion into the model's generations, known as the Prompt Boundary Problem (PBP)<n>We present an inference-time method to convert any autore LM with a BPE tokenizer into a character-level or byte-level LM.<n>Our method efficiently solves the PBP and is also able to unify the vocabularies of language models with different tokenizers.
  arXiv  Detail & Related papers  (2025-06-17T02:37:04Z)
• Structural Entropy Guided Probabilistic Coding [52.01765333755793]
  We propose a novel structural entropy-guided probabilistic coding model, named SEPC.<n>We incorporate the relationship between latent variables into the optimization by proposing a structural entropy regularization loss.<n> Experimental results across 12 natural language understanding tasks, including both classification and regression tasks, demonstrate the superior performance of SEPC.
  arXiv  Detail & Related papers  (2024-12-12T00:37:53Z)
• Exact Byte-Level Probabilities from Tokenized Language Models for FIM-Tasks and Model Ensembles [23.134664392314264]
  Tokenization is associated with many poorly understood shortcomings in language models (LMs)<n>This work studies how tokenization impacts model performance by analyzing and comparing models with their byte-level counterparts.<n>We introduce the Byte-Token Representation Lemma, a framework that establishes a mapping between the learned token distribution and its equivalent byte-level distribution.
  arXiv  Detail & Related papers  (2024-10-11T23:30:42Z)
• Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion [61.03681839276652]
  Diffusion Forcing is a new training paradigm where a diffusion model is trained to denoise a set of tokens with independent per-token noise levels.<n>We apply Diffusion Forcing to sequence generative modeling by training a causal next-token prediction model to generate one or several future tokens.
  arXiv  Detail & Related papers  (2024-07-01T15:43:25Z)
• Amortizing intractable inference in large language models [56.92471123778389]
  We use amortized Bayesian inference to sample from intractable posterior distributions. We empirically demonstrate that this distribution-matching paradigm of LLM fine-tuning can serve as an effective alternative to maximum-likelihood training. As an important application, we interpret chain-of-thought reasoning as a latent variable modeling problem.
  arXiv  Detail & Related papers  (2023-10-06T16:36:08Z)
• SeqDiffuSeq: Text Diffusion with Encoder-Decoder Transformers [50.90457644954857]
  In this work, we apply diffusion models to approach sequence-to-sequence text generation. We propose SeqDiffuSeq, a text diffusion model for sequence-to-sequence generation. Experiment results illustrate the good performance on sequence-to-sequence generation in terms of text quality and inference time.
  arXiv  Detail & Related papers  (2022-12-20T15:16:24Z)
• SparseGAN: Sparse Generative Adversarial Network for Text Generation [8.634962333084724]
  We propose a SparseGAN that generates semantic-interpretable, but sparse sentence representations as inputs to the discriminator. With such semantic-rich representations, we not only reduce unnecessary noises for efficient adversarial training, but also make the entire training process fully differentiable.
  arXiv  Detail & Related papers  (2021-03-22T04:44:43Z)
• Improve Variational Autoencoder for Text Generationwith Discrete Latent Bottleneck [52.08901549360262]
  Variational autoencoders (VAEs) are essential tools in end-to-end representation learning. VAEs tend to ignore latent variables with a strong auto-regressive decoder. We propose a principled approach to enforce an implicit latent feature matching in a more compact latent space.
  arXiv  Detail & Related papers  (2020-04-22T14:41:37Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.