[논문 리뷰] Blockwise Self-Attention for Long Document Understanding

Jiezhong Qiu, Hao Ma외 4명|arXiv (Cornell University)|2019. 11. 07.

Topic Modeling참고 문헌 44인용 수 16

한 줄 요약

BlockBERT는 BERT의 자기주의 메커니즘에 희소 블록 구조를 도입하여 전체 시퀀스가 아닌 국소 블록으로 주의를 제한함으로써 메모리 및 계산 비용을 감소시킨다. RoBERTa 대비 18.7–36.1% 낮은 메모리 사용량, 12.0–25.1% 빠른 훈련 속도, 27.8% 빠른 추론 속도를 기록하며, 장문 문서 작업에서 동등하거나 향상된 최종 성능을 달성한다.

ABSTRACT

We present BlockBERT, a lightweight and efficient BERT model for better modeling long-distance dependencies. Our model extends BERT by introducing sparse block structures into the attention matrix to reduce both memory consumption and training/inference time, which also enables attention heads to capture either short- or long-range contextual information. We conduct experiments on language model pre-training and several benchmark question answering datasets with various paragraph lengths. BlockBERT uses 18.7-36.1% less memory and 12.0-25.1% less time to learn the model. During testing, BlockBERT saves 27.8% inference time, while having comparable and sometimes better prediction accuracy, compared to an advanced BERT-based model, RoBERTa.

연구 동기 및 목표

• 장문 시퀀스에서 BERT를 훈련하고 추론할 때 발생하는 높은 메모리 및 계산 비용을 해결하여 모델의 확장성과 구현 가능성을 향상시키기 위해.
• 시퀀스 길이에 따라 제곱 복잡도를 가지는 도트곱 자기주의가 BERT에서 주요 메모리 병목 현상임을 규명하기 위해.
• 장기적 맥락 작업에서 강력한 성능을 유지하면서도 메모리 및 추론 시간을 크게 줄이는 경량 BERT 변형을 설계하기 위해.
• 모델 정확도를 훼손하지 않으면서 주의 행렬의 구조적 희소성으로 장거리 의존성을 효율적으로 모델링하기 위해.
• 복잡한 희소 주의 메커니즘에 의존하지 않고, 커스텀 커널이나 제한된 하드웨어 지원이 필요 없는 간단하고 구현 가능한 대안을 제공하기 위해.

제안 방법

• 시퀀스를 고정 크기의 블록으로 나누고, 자기주의를 블록 내부 및 선택된 블록 간 연결에 국한함으로써 블록 희소 주의를 도입하기 위해.
• BERT의 전체 밀도 주의 행렬을 블록 희소 행렬 구조로 대체하여, 희소성 패턴에 비례해 FLOPs와 메모리 사용량을 감소시키기 위해.
• 다중 헤드 주의를 유지하되, 각 헤드가 오직 정의된 블록 내부 또는 블록 간에만 주의를 기울이도록 제한하여 국소 및 장거리 모델링을 모두 가능하게 하기 위해.
• 표준 딥 러닝 프레임워크와 호환되며 커스텀 CUDA 커널이 필요 없는 간단하고 미분 가능한 희소성 패턴을 사용하기 위해.
• 원본 BERT 아키텍처와 사전 훈련 목표를 그대로 유지하면서, 효율성을 향상시키기 위해 주의 메커니즘만 수정하기 위해.
• 모든 트랜스포머 레이어에 블록 희소 주의를 적용하여 기존 미세조정 파이프라인과의 종단 간 호환성을 보장하기 위해.

실험 결과

연구 질문

• RQ1 구조적 희소성이 주의 행렬에 도입될 경우, 장기 맥락 작업에서 성능 저하 없이 BERT의 메모리 및 계산 비용을 줄일 수 있는가?
• RQ2 장문 문서에서 블록 희소 주의는 밀도 주의보다 다운스트림 작업 정확도에서 어떻게 비교되는가?
• RQ3 블록 희소 주의는 모델 용량을 유지하면서 훈련 및 추론 시간을 얼마나 줄일 수 있는가?
• RQ4 블록 희소 설계는 주의 헤드가 의도한 바와 같이 짧은 거리와 장거리 의존성을 효과적으로 포착할 수 있는가?
• RQ5 실제 구현 환경에서 RoBERTa와 같은 최신 BERT 변형 대비 BlockBERT의 효율성 향상 정도는 어떠한가?

주요 결과

• BlockBERT는 RoBERTa 대비 훈련 시 18.7–36.1% 낮은 메모리 소비를 기록하며, 주로 주의 행렬의 희소화 덕분이다.
• 더 적은 FLOPs와 낮은 메모리 압박 덕분에 더 큰 배치 크기를 허용함으로써 훈련 시간이 12.0–25.1% 감소한다.
• RoBERTa 대비 추론 시간이 27.8% 감소하여 실시간 또는 자원 제약이 있는 응용 분야에 더 적합하다.
• 다양한 단락 길이를 가진 질문-응답 벤치마크에서 BlockBERT는 RoBERTa와 동등하거나 더 높은 성능을 달성한다.
• 블록 희소 주의 메커니즘이 주의 헤드가 국소적 및 장거리 의존성을 효과적으로 모델링하는 데 성공했다.
• 이 방법은 단순하며 표준 딥 러닝 프레임워크와 호환되며, 커스텀 커널이나 복잡한 희소성 관리가 필요 없다.