TensorBLEU: Implementação Vetorizada Baseada em GPU para Cálculo de Pontuação BLEU por Sentença Durante o Treinamento

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

October 7, 2025

Autores: Adam Filipek

cs.AI

Resumo

Modelos modernos de processamento de linguagem natural atingiram uma escala sem precedentes, mas as ferramentas para sua avaliação frequentemente permanecem um gargalo computacional, limitando o ritmo da pesquisa. Isso é particularmente crítico para métricas de avaliação durante o treinamento, como sinais de recompensa por frase em Aprendizado por Reforço, que devem operar de forma eficiente em lotes de IDs de tokens diretamente na GPU. Neste artigo, apresentamos o TensorBLEU, uma nova implementação da métrica BLEU projetada desde o início para esse caso de uso específico. Nossa abordagem é totalmente vetorizada para computação acelerada por GPU, por frase, dentro do PyTorch, e introduz um mecanismo de contagem eficiente em termos de memória. Ao criar um dicionário compacto e específico para o lote de n-grams usando torch.unique, nosso método evita os custos proibitivos de memória da vetorização tradicional baseada em hashing, tornando-o prático para modelos de vocabulário grande. Avaliamos o TensorBLEU em comparação com o NLTK, a biblioteca padrão para cálculo de BLEU baseado em IDs de tokens na CPU. Experimentos mostram que o TensorBLEU oferece acelerações de mais de 13x em GPUs de consumo (NVIDIA T4) e excede 40x em hardware de classe datacenter (NVIDIA A100). Esse desempenho transforma um gargalo significativo em uma parte insignificante do loop de treinamento. Ao definir claramente seu papel como um "BLEU de IDs de Tokens" para fins de desenvolvimento e ao disponibilizar nossa implementação em código aberto, fornecemos uma ferramenta poderosa para acelerar a pesquisa em áreas como o ajuste fino de modelos baseados em 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.