TensorBLEU: Vektorisierte GPU-basierte BLEU-Score-Implementierung für die Bewertung pro Satz während des Trainings

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

October 7, 2025

papers.authors: Adam Filipek

cs.AI

papers.abstract

Moderne Modelle der natürlichen Sprachverarbeitung haben ein bisher unerreichtes Maß an Skalierung erreicht, doch die Werkzeuge zu ihrer Bewertung stellen oft einen rechenintensiven Engpass dar, der das Forschungstempo begrenzt. Dies ist besonders gravierend bei Metriken zur Bewertung während des Trainings, wie zum Beispiel belohnungsbasierte Signale pro Satz im Reinforcement Learning, die effizient auf Batches von Token-IDs direkt auf der GPU operieren müssen. In diesem Artikel stellen wir TensorBLEU vor, eine neuartige Implementierung der BLEU-Metrik, die speziell für diesen Anwendungsfall entwickelt wurde. Unser Ansatz ist vollständig vektorisiert für GPU-beschleunigte, satzweise Berechnungen innerhalb von PyTorch und führt einen speicher-effizienten Zählmechanismus ein. Durch die Erstellung eines kompakten, batch-spezifischen Wörterbuchs von n-Grammen mittels torch.unique vermeidet unsere Methode die prohibitiv hohen Speicherkosten traditioneller, hash-basierter Vektorisierung, was sie für Modelle mit großem Vokabular praktikabel macht. Wir vergleichen TensorBLEU mit NLTK, der Standardbibliothek für die BLEU-Berechnung auf Basis von Token-IDs auf der CPU. Experimente zeigen, dass TensorBLEU Beschleunigungen von über 13x auf Consumer-GPUs (NVIDIA T4) und mehr als 40x auf Data-Center-Hardware (NVIDIA A100) bietet. Diese Leistung wandelt einen signifikanten Engpass in einen vernachlässigbaren Teil des Trainingszyklus um. Indem wir seine Rolle als „Token-ID-BLEU“ für Entwicklungszwecke klar definieren und unsere Implementierung Open-Source zur Verfügung stellen, bieten wir ein leistungsstarkes Werkzeug zur Beschleunigung der Forschung in Bereichen wie der feinabstimmenden Modelloptimierung mittels Reinforcement Learning.

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.