Zonkey: A Hierarchical Diffusion Language Model with Differentiable Tokenization and Probabilistic Attention
作者: Alon Rozental
分类: cs.CL
发布日期: 2026-01-29
💡 一句话要点
Zonkey:提出一种可微分分词和概率注意力机制的层级扩散语言模型,实现端到端优化。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 扩散模型 可微分分词 概率注意力 自然语言生成 端到端学习
📋 核心要点
- 现有LLM受限于不可微的固定分词器(如BPE),无法进行端到端优化,且难以适应噪声或特定领域数据。
- Zonkey提出一种层级扩散模型,核心是可微分词器和概率注意力机制,实现从字符到文档的端到端训练。
- Zonkey在Wikipedia上训练,能从噪声生成连贯文本,展示了涌现的层次结构,并与数据分布良好对齐。
📝 摘要(中文)
大型语言模型(LLMs)彻底改变了自然语言处理,但仍受限于固定的、不可微的分词器,如字节对编码(BPE),这阻碍了端到端优化以及对噪声或特定领域数据的适应性。我们提出了Zonkey,一种层级扩散模型,通过一个从原始字符到文档级表示的完全可训练的流程来解决这些限制。其核心是一个可微的分词器(Segment Splitter),它学习概率性的序列起始(BOS)决策,从而实现自适应分割,这些分割在没有明确监督的情况下呈现出语言学意义(例如,空格处的单词边界,句点处的句子开始)。这种可微性是通过我们新颖的概率注意力机制实现的,该机制结合了位置特定的存在概率来模拟理论上无限序列上的软掩码,同时保留梯度。序列以概率方式衰减,而不是依赖于序列结束标记,从而支持可变长度的输出。层级级别将序列压缩为更高的抽象(例如,字符n-gram到类似单词的向量,然后是类似句子的向量),并通过我们的去噪扩散混合模型(DDMM)进行重建,以在潜在空间中实现稳定和高效的去噪。一个Stitcher确保了段之间的重叠不变性。Zonkey在Wikipedia上进行端到端训练,从噪声中生成连贯的、可变长度的文本,展示了涌现的层次结构,并与基于熵的可学习分词器相比,在质量上与数据分布具有良好的对齐。我们的方法朝着完全基于梯度的LLM迈进了一步,具有更好的领域适应性和可扩展生成潜力。我们发布了用于训练和复现我们实验的源代码。
🔬 方法详解
问题定义:现有大型语言模型依赖于固定的、不可微的分词器,例如Byte Pair Encoding (BPE)。这些分词器阻碍了模型的端到端优化,并且难以适应噪声数据或特定领域的数据。因此,如何设计一种可微分的分词方法,使得模型能够更好地学习语言的内在结构,并提升在各种数据上的泛化能力,是本文要解决的核心问题。
核心思路:Zonkey的核心思路是使用扩散模型,并引入可微分的分词器和概率注意力机制,从而实现从原始字符到文档级表示的完全可训练的流程。通过可微分的分词器,模型可以自适应地学习token的边界,而无需人工标注。概率注意力机制则允许模型处理变长序列,并保留梯度信息,从而实现端到端的优化。
技术框架:Zonkey的技术框架包含以下几个主要模块:1) 可微分分词器 (Segment Splitter):学习概率性的序列起始 (BOS) 决策,实现自适应分割。2) 概率注意力机制 (Probabilistic Attention):结合位置特定的存在概率来模拟软掩码,处理变长序列并保留梯度。3) 层级结构 (Hierarchical Levels):将序列压缩为更高的抽象表示,例如字符 n-gram 到类似单词的向量,然后是类似句子的向量。4) 去噪扩散混合模型 (DDMM):在潜在空间中进行稳定和高效的去噪,用于序列的重建。5) Stitcher:确保段之间的重叠不变性。
关键创新:Zonkey最重要的技术创新点在于其可微分的分词器和概率注意力机制。传统的分词器是不可微的,这使得模型无法进行端到端的优化。Zonkey通过学习概率性的序列起始决策,实现了可微分的分词,从而允许模型更好地学习语言的内在结构。概率注意力机制则允许模型处理变长序列,并保留梯度信息,这使得模型能够更好地处理复杂的语言结构。
关键设计:可微分分词器通过学习每个字符作为序列起始的概率来实现。概率注意力机制通过引入位置特定的存在概率来模拟软掩码,从而处理变长序列。去噪扩散混合模型 (DDMM) 的具体实现细节(例如,噪声 schedule、网络结构等)未知。损失函数的设计也未知,但应该包含重建损失和分词的正则化项。
🖼️ 关键图片
📊 实验亮点
Zonkey在Wikipedia数据集上进行了端到端训练,并生成了连贯的、可变长度的文本。实验结果表明,Zonkey能够涌现出层次结构,并且在质量上与数据分布具有良好的对齐。与基于熵的可学习分词器相比,Zonkey在生成文本的质量上有所提升,但具体的性能数据和提升幅度未知。
🎯 应用场景
Zonkey的潜在应用领域包括自然语言生成、机器翻译、文本摘要等。其可微分的分词方法和概率注意力机制可以提升模型在噪声数据和特定领域数据上的性能。该研究的实际价值在于提供了一种新的端到端训练大型语言模型的方法,未来可能促进更高效、更灵活的语言模型的发展。
📄 摘要(原文)
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.