TensorBLEU: 훈련 중 문장별 평가를 위한 벡터화된 GPU 기반 BLEU 점수 구현

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

October 7, 2025

저자: Adam Filipek

cs.AI

초록

현대 자연어 처리 모델은 전례 없는 규모를 달성했으나, 이들의 평가 도구는 종종 계산적 병목 현상으로 남아 연구 속도를 제한하고 있다. 이는 특히 강화 학습에서의 문장별 보상 신호와 같은 훈련 중 평가 지표에서 두드러지는데, 이러한 지표는 GPU에서 직접 토큰 ID 배치에 대해 효율적으로 작동해야 한다. 본 논문에서는 이러한 특정 사용 사례를 위해 처음부터 설계된 BLEU 지표의 새로운 구현체인 TensorBLEU를 소개한다. 우리의 접근 방식은 PyTorch 내에서 GPU 가속화된 문장별 계산을 위해 완전히 벡터화되었으며, 메모리 효율적인 카운팅 메커니즘을 도입했다. torch.unique를 사용하여 n-그램의 컴팩트하고 배치별 사전을 생성함으로써, 우리의 방법은 전통적인 해싱 기반 벡터화의 과도한 메모리 비용을 피해 대규모 어휘 모델에 실용적으로 적용할 수 있다. 우리는 TensorBLEU를 CPU 기반 토큰 ID BLEU 계산을 위한 표준 라이브러리인 NLTK와 비교 평가했다. 실험 결과, TensorBLEU는 소비자 등급 GPU(NVIDIA T4)에서 13배 이상, 데이터 센터 등급 하드웨어(NVIDIA A100)에서는 40배 이상의 속도 향상을 제공하는 것으로 나타났다. 이 성능은 상당한 병목 현상을 훈련 루프에서 무시할 수 있는 수준으로 전환한다. 개발 목적을 위한 "토큰 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.