TensorBLEU: Векторизованная реализация оценки BLEU на GPU для поэтапной оценки предложений в процессе обучения

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

October 7, 2025

Авторы: Adam Filipek

cs.AI

Аннотация

Современные модели обработки естественного языка достигли беспрецедентного масштаба, однако инструменты для их оценки часто остаются вычислительным узким местом, ограничивая темпы исследований. Это особенно актуально для метрик, используемых во время обучения, таких как сигналы вознаграждения на уровне предложений в обучении с подкреплением, которые должны эффективно работать с пакетами идентификаторов токенов непосредственно на GPU. В данной статье мы представляем TensorBLEU — новую реализацию метрики BLEU, разработанную с нуля для этого конкретного случая использования. Наш подход полностью векторизован для ускоренных вычислений на GPU на уровне предложений в рамках PyTorch и включает в себя механизм подсчета, оптимизированный по использованию памяти. Создавая компактный, специфичный для пакета словарь n-грамм с использованием torch.unique, наш метод избегает непомерных затрат памяти, характерных для традиционной векторизации на основе хэширования, что делает его практичным для моделей с большим словарным запасом. Мы сравниваем TensorBLEU с NLTK — стандартной библиотекой для расчета BLEU на основе идентификаторов токенов на CPU. Эксперименты показывают, что TensorBLEU обеспечивает ускорение более чем в 13 раз на потребительских GPU (NVIDIA T4) и превышает 40 раз на оборудовании класса дата-центров (NVIDIA A100). Такая производительность превращает значительное узкое место в незначительную часть цикла обучения. Четко определив свою роль как "Token-ID BLEU" для целей разработки и открыв исходный код нашей реализации, мы предоставляем мощный инструмент для ускорения исследований в таких областях, как тонкая настройка моделей на основе обучения с подкреплением.

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.