TensorBLEU：用于训练中逐句评估的向量化 GPU BLEU 分数实现

现代自然语言处理模型已达到前所未有的规模，然而其评估工具往往仍是计算瓶颈，限制了研究进度。这对于训练中评估指标尤为突出，例如强化学习中的逐句奖励信号，必须直接在 GPU 上高效处理 token ID 批次。在本文中，我们介绍了 TensorBLEU，一种专为这一特定用例从头构建的新型 BLEU 指标实现。我们的方法完全向量化，用于在 PyTorch 中进行 GPU 加速的逐句计算，并引入了一种内存高效的计数机制。通过使用 \texttt{torch.unique} 创建紧凑的、特定于批次的 n-gram 字典，我们的方法避免了传统基于哈希的向量化带来的高昂内存成本，使其适用于大词汇量模型。我们将 TensorBLEU 与在 CPU 上基于 token ID 计算 BLEU 的标准库 NLTK 进行了基准对比。实验表明，TensorBLEU 在消费级 GPU（NVIDIA T4）上提供了超过 13 倍的加速，在数据中心级硬件（NVIDIA A100）上加速超过 40 倍。这一性能将一个重大瓶颈转化为训练循环中可忽略不计的部分。通过明确界定其作为开发用途“Token-ID BLEU”的角色并开源我们的实现，我们为加速基于 RL 的模型微调等领域的研究提供了强大的工具。