TensorBLEU: Implementación Vectorizada Basada en GPU para la Evaluación de la Puntuación BLEU por Oración Durante el Entrenamiento

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

October 7, 2025

Autores: Adam Filipek

cs.AI

Resumen

Los modelos modernos de procesamiento del lenguaje natural han alcanzado una escala sin precedentes, sin embargo, las herramientas para su evaluación a menudo siguen siendo un cuello de botella computacional, limitando el ritmo de la investigación. Esto es particularmente crítico para las métricas de evaluación durante el entrenamiento, como las señales de recompensa por oración en el Aprendizaje por Refuerzo, que deben operar de manera eficiente en lotes de identificadores de tokens directamente en la GPU. En este artículo, presentamos TensorBLEU, una implementación novedosa de la métrica BLEU diseñada desde cero para este caso de uso específico. Nuestro enfoque está completamente vectorizado para el cálculo por oración acelerado por GPU dentro de PyTorch e introduce un mecanismo de conteo eficiente en memoria. Al crear un diccionario compacto y específico para el lote de n-gramas utilizando torch.unique, nuestro método evita los costos prohibitivos de memoria de la vectorización tradicional basada en hashing, haciéndolo práctico para modelos de vocabulario grande. Evaluamos TensorBLEU en comparación con NLTK, la biblioteca estándar para el cálculo de BLEU basado en identificadores de tokens en la CPU. Los experimentos muestran que TensorBLEU ofrece aceleraciones de más de 13x en GPUs de consumo (NVIDIA T4) y supera 40x en hardware de clase centro de datos (NVIDIA A100). Este rendimiento transforma un cuello de botella significativo en una parte insignificante del ciclo de entrenamiento. Al definir claramente su papel como un "BLEU de identificadores de tokens" para fines de desarrollo y al liberar nuestra implementación como código abierto, proporcionamos una herramienta poderosa para acelerar la investigación en áreas como el ajuste fino de modelos basados en Aprendizaje por Refuerzo.

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.