DNATokenizer: A GPU-First Byte-to-Identifier Tokenizer for High-Throughput DNA Language Models
作者: Eliatan Niktab, Hardip Patel
分类: q-bio.GN, cs.LG
发布日期: 2026-01-09
💡 一句话要点
提出DNATokenizer,一种GPU优先的字节到标识符分词器,用于高吞吐量DNA语言模型。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 基因组语言模型 分词器 GPU加速 生物信息学 高吞吐量计算
📋 核心要点
- 现有基因组语言模型的分词方法,如重叠k-mer、单核苷酸和BPE,在信息泄露、序列长度和边界敏感性等方面存在局限性。
- DNATok通过GPU优先的字节查找表和重叠H2D/计算流水线,实现了高性能的分词,并且与具体的分词算法和词汇表无关。
- DNATok的编码吞吐量比Hugging Face基线高84-95倍,H2D吞吐量提升高达1.9倍,端到端流式传输速率达到1.27-1.84e8 tokens/s。
📝 摘要(中文)
分词位于高吞吐量基因组输入和GPU计算的边界,在算法设计和系统吞吐量方面都带来了挑战。重叠k-mer分词可能在掩码语言建模(MLM)下引入信息泄露,并可能降低下游准确性。单核苷酸分词避免了泄露并保留了每个碱基的保真度,但它大大增加了基于注意力的架构的序列长度。非重叠k-mers和字节对编码(BPE)提供了压缩并避免了泄露,但代价是边界敏感性或降低了可解释性。经验表明,分词的选择与模型架构和任务需求密切相关。然而,在系统层面,一旦输入达到数十亿个碱基,标准字符串分词器和主机绑定的词汇查找就会占据挂钟时间,而与分词算法无关。我们提出了DNATok,一种高性能、GPU优先的分词系统,它用基于字节查找表(LUT)的标识符流和使用固定内存和架构并行性的重叠主机到设备(H2D)/计算流水线取代了通用字符串处理。DNATok与词汇无关:它加速了单核苷酸、非重叠k-mer和BPE分词,并作为基因组基础模型下的一个即插即用系统层集成。DNATok实现了比优化的Hugging Face基线高84-95倍的编码吞吐量,以及高达1.9倍的H2D吞吐量。端到端流式传输达到1.27-1.84e8 tokens/s,具体取决于配置,有效地消除了分词作为生产规模训练和推理的瓶颈。
🔬 方法详解
问题定义:现有基因组语言模型的分词方法存在效率瓶颈,标准字符串分词器和主机绑定的词汇查找在处理大规模基因组数据时会显著降低整体性能。此外,不同的分词算法(如k-mer、单核苷酸、BPE)各有优缺点,需要一种通用的加速方案。
核心思路:DNATok的核心思路是将分词过程尽可能地卸载到GPU上进行,利用GPU的并行计算能力加速字节到标识符的转换。通过优化数据传输和计算流程,减少主机和设备之间的通信开销,从而提高整体分词吞吐量。
技术框架:DNATok的整体架构包含以下几个主要阶段:1) 输入数据读取;2) 基于字节查找表(LUT)的标识符流生成;3) 主机到设备(H2D)的数据传输,采用pinned memory和重叠的H2D/计算流水线;4) GPU上的分词计算。该框架支持单核苷酸、非重叠k-mer和BPE等多种分词算法。
关键创新:DNATok的关键创新在于其GPU优先的设计理念和字节查找表的使用。通过将大部分分词计算转移到GPU上,充分利用了GPU的并行计算能力。字节查找表避免了复杂的字符串处理操作,提高了分词效率。重叠的H2D/计算流水线进一步减少了数据传输的瓶颈。
关键设计:DNATok的关键设计包括:1) 使用pinned memory来加速主机和设备之间的数据传输;2) 设计高效的字节查找表,将每个字节映射到对应的标识符;3) 实现重叠的H2D/计算流水线,使得数据传输和计算可以并行进行;4) 针对不同的分词算法,优化GPU上的计算kernel。
📊 实验亮点
DNATok在基因组数据分词方面取得了显著的性能提升。与优化的Hugging Face基线相比,DNATok实现了84-95倍更高的编码吞吐量,以及高达1.9倍的H2D吞吐量。端到端流式传输速率达到1.27-1.84e8 tokens/s,有效地消除了分词作为生产规模训练和推理的瓶颈。这些结果表明DNATok是一种高效且可扩展的基因组数据分词解决方案。
🎯 应用场景
DNATok可应用于各种需要处理大规模基因组数据的场景,例如基因组序列分析、基因组语言模型训练和推理、药物发现等。通过提高分词效率,DNATok可以加速基因组研究的进程,并降低计算成本。未来,DNATok可以进一步扩展到其他生物序列数据处理领域,例如蛋白质序列分析。
📄 摘要(原文)
Tokenization sits at the boundary between high-throughput genomic input and GPU compute, posing challenges in both algorithm design and system throughput. Overlapping k-mer tokenization can introduce information leakage under masked language modeling (MLM) and may degrade downstream accuracy. Single-nucleotide tokenization avoids leakage and preserves per-base fidelity, but it greatly increases sequence length for attention-based architectures. Non-overlapping k-mers and byte-pair encoding (BPE) provide compression and avoid leakage, at the cost of boundary sensitivity or reduced interpretability. Empirically, the choice of tokenization interacts strongly with model architecture and task requirements. At the system level, however, standard string tokenizers and host-bound vocabulary lookups dominate wall-clock time once inputs reach billions of bases, regardless of the tokenization algorithm. We present DNATok, a high-performance, GPU-first tokenization system that replaces general-purpose string processing with byte lookup table (LUT)-based identifier streaming and an overlapped host-to-device (H2D)/compute pipeline using pinned memory and architectural parallelism. DNATok is vocabulary-agnostic: it accelerates single-nucleotide, non-overlapping k-mer, and BPE tokenization, and integrates as a drop-in systems layer beneath genomic foundation models. DNATok achieves 84-95x higher encoding throughput than optimized Hugging Face baselines and up to 1.9x higher H2D throughput. End-to-end streaming reaches 1.27-1.84e8 tokens/s depending on configuration, effectively removing tokenization as a bottleneck for production-scale training and inference.