TensorBLEU: Implementazione Vettoriale Basata su GPU del Punteggio BLEU per la Valutazione Per-Frase Durante l'Addestramento

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

October 7, 2025

Autori: Adam Filipek

cs.AI

Abstract

I moderni modelli di elaborazione del linguaggio naturale hanno raggiunto una scala senza precedenti, tuttavia gli strumenti per la loro valutazione spesso rimangono un collo di bottiglia computazionale, limitando il ritmo della ricerca. Questo è particolarmente evidente per le metriche di valutazione durante l'addestramento, come i segnali di ricompensa per frase nel Reinforcement Learning, che devono operare in modo efficiente su batch di ID di token direttamente sulla GPU. In questo articolo, introduciamo TensorBLEU, una nuova implementazione della metrica BLEU progettata da zero per questo specifico caso d'uso. Il nostro approccio è completamente vettorizzato per il calcolo per frase accelerato da GPU all'interno di PyTorch e introduce un meccanismo di conteggio efficiente in termini di memoria. Creando un dizionario compatto e specifico per batch di n-grammi utilizzando torch.unique, il nostro metodo evita i costi di memoria proibitivi della vettorizzazione tradizionale basata su hashing, rendendolo pratico per modelli con vocabolari estesi. Abbiamo confrontato TensorBLEU con NLTK, la libreria standard per il calcolo di BLEU basato su ID di token sulla CPU. Gli esperimenti mostrano che TensorBLEU offre accelerazioni superiori a 13x su GPU di livello consumer (NVIDIA T4) e oltre 40x su hardware di classe data center (NVIDIA A100). Questa prestazione trasforma un significativo collo di bottiglia in una parte trascurabile del ciclo di addestramento. Definendo chiaramente il suo ruolo come "Token-ID BLEU" per scopi di sviluppo e rendendo open-source la nostra implementazione, forniamo uno strumento potente per accelerare la ricerca in aree come il fine-tuning di modelli basato su RL.

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.