TensorBLEU : Implémentation vectorisée basée sur GPU du score BLEU pour l'évaluation par phrase pendant l'entraînement

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

October 7, 2025

papers.authors: Adam Filipek

cs.AI

papers.abstract

Les modèles modernes de traitement du langage naturel ont atteint une échelle sans précédent, mais les outils pour leur évaluation restent souvent un goulot d'étranglement computationnel, limitant le rythme de la recherche. Ce problème est particulièrement aigu pour les métriques d'évaluation en cours d'entraînement, telles que les signaux de récompense par phrase dans l'apprentissage par renforcement, qui doivent fonctionner efficacement sur des lots d'identifiants de tokens directement sur le GPU. Dans cet article, nous présentons TensorBLEU, une nouvelle implémentation de la métrique BLEU conçue spécifiquement pour ce cas d'utilisation. Notre approche est entièrement vectorisée pour un calcul accéléré par GPU, phrase par phrase, dans PyTorch, et introduit un mécanisme de comptage économe en mémoire. En créant un dictionnaire compact de n-grams spécifique à chaque lot à l'aide de torch.unique, notre méthode évite les coûts mémoire prohibitifs de la vectorisation traditionnelle basée sur le hachage, la rendant pratique pour les modèles à vocabulaire étendu. Nous comparons TensorBLEU à NLTK, la bibliothèque standard pour le calcul de BLEU basé sur des identifiants de tokens sur le CPU. Les expériences montrent que TensorBLEU offre des accélérations de plus de 13x sur des GPU grand public (NVIDIA T4) et dépassant 40x sur du matériel de classe centre de données (NVIDIA A100). Cette performance transforme un goulot d'étranglement significatif en une partie négligeable de la boucle d'entraînement. En définissant clairement son rôle en tant que "BLEU sur identifiants de tokens" à des fins de développement et en ouvrant notre implémentation en open source, nous fournissons un outil puissant pour accélérer la recherche dans des domaines comme le réglage fin de modèles basé sur l'apprentissage par renforcement.

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 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.