arXiv:2510.05485cs.CLcs.LG2025-10

用GPU加速句子级BLEU计算,让模型训练评估快13倍以上。

TensorBLEU: Vectorized GPU-based BLEU Score Implementation for Per-Sentence In-Training Evaluation

  • 基于PyTorch全向量化设计,直接在GPU上处理词元ID批数据
  • 采用紧凑字典机制避免内存爆炸,支持大词汇量模型
  • 适合强化学习微调等需高频评估的场景,已开源

现代自然语言处理模型规模空前,但评估工具常成为计算瓶颈,尤其在训练中需要逐句奖励信号时。本文提出TensorBLEU,一种专为该场景设计的新型BLEU实现。方法完全向量化,可在PyTorch中直接于GPU上进行批量词元ID的句子级计算,并通过torch.unique构建紧凑的、批次专用的n-gram字典,避免传统哈希向量化带来的高昂内存开销。在NVIDIA T4(消费级)和A100(数据中心级)上分别实现超过13倍和40倍的加速。实验表明,该方法将原本显著的瓶颈转化为训练循环中的可忽略部分。通过明确其作为开发用途的“词元ID BLEU”角色并开源实现,为基于强化学习的模型微调等研究提供了高效工具。

原文摘要 · Abstract (English)

Modern natural language processing models have achieved unprecedented scale, yet the tools for their evaluation often remain a computational bottleneck, limiting the pace of research. This is particularly acute for in-training evaluation metrics, such as per-sentence reward signals in Reinforcement Learning, which must operate efficiently on batches of token IDs directly on the GPU. In this paper, we introduce TensorBLEU, a novel implementation of the BLEU metric designed from the ground up for this specific use case. Our approach is fully vectorized for GPU-accelerated, per-sentence computation within PyTorch and introduces a memory-efficient counting mechanism. By creating a compact, batch-specific dictionary of n-grams using \texttt{torch.unique}, our method avoids the prohibitive memory costs of traditional hashing-based vectorization, making it practical for large-vocabulary models. We benchmark TensorBLEU against NLTK, the standard library for token-ID-based BLEU calculation on the CPU. Experiments show that TensorBLEU provides speedups of over 13x on consumer-grade GPUs (NVIDIA T4) and exceeding 40x on data-center-class hardware (NVIDIA A100). This performance transforms a significant bottleneck into a negligible part of the training loop. By clearly defining its role as a "Token-ID BLEU" for development purposes and open-sourcing our implementation, we provide a powerful tool for accelerating research in areas like RL-based model fine-tuning.

BLEUGPU加速训练评估PyTorch

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。