RED-PIM: Processing-in-Memory를 사용하여 Transformer의 데이터 이동 감소
요약
RED-PIM은 Processing-in-Memory(PIM) 기술을 활용하여 Transformer 모델의 어텐션 연산 시 발생하는 데이터 이동 문제를 해결하는 알고리즘-아키텍처 공동 설계 연구입니다. 뱅크 간 통신을 최소화하고 중간 행렬 크기를 축소하여 추론 성능과 확장성을 크게 개선했습니다.
핵심 포인트
- PIM 기반 Transformer의 뱅크 간 데이터 이동을 O(N^2)에서 O(N)으로 감소
- 중간 어텐션 행렬을 N x N에서 d x d로 축소하여 지연 시간 단축
- 기존 PIM 구현 대비 추론 시간을 평균 66.42% 단축
- 긴 시퀀스 데이터에서 성능 개선 효과가 더욱 두드러짐
- 정확도를 유지하거나 향상시키면서 효율적인 추론 가능
Transformers는 자연어 처리 (NLP), 컴퓨터 비전 (Computer Vision), 웹 검색, DNA 서열 분석을 포함한 많은 도메인에서 널리 사용됩니다. 이러한 광범위한 적용 가능성을 고려할 때, transformer 모델의 성능을 향상시키는 것은 매우 중요합니다. 그러나 어텐션 (Attention) 연산 중에 프로세싱 유닛 (Processing Units)과 메모리 사이에서 발생하는 대량의 데이터 이동은 효율성을 크게 제한합니다. Processing-In-Memory (PIM)는 메모리 내부에서 직접 연산을 수행함으로써 이 문제를 완화합니다. 기존 연구들이 PIM 기반의 transformer 구현을 제안해 왔으나, 이들은 비용이 많이 드는 뱅크 간 통신 (Inter-bank communication) 문제로 어려움을 겪으며, 메모리 뱅크 (Memory banks)의 제한된 용량으로 인해 확장성 확보에 어려움이 있습니다. 결과적으로, 어텐션 관련 데이터를 여러 뱅크에 나누어 배치해야 하며, 이는 PIM의 잠재적 이점을 감소시킵니다. 본 연구에서는 뱅크 간 데이터 이동을 $O(N^2)$에서 $O(N)$으로 최소화하고, 중간 어텐션 행렬 (Intermediate attention matrices)을 $N \times N$에서 $d \times d$로 축소하여 어텐션 지연 시간 (Attention latency)을 줄이는 알고리즘-아키텍처 공동 설계 (Algorithm-architecture co-design)인 RED-PIM을 제안합니다. 행렬 연산을 재구성하고, 연산을 로컬에서 수행하며, 최적화된 데이터 전송 전략을 채택함으로써, RED-PIM은 연산 비용과 상호 연결 트래픽 (Interconnect traffic)을 크게 줄입니다. 베이스라인 PIM 구현과 비교했을 때, RED-PIM은 추론 시간 (Inference time)을 16.05%에서 99.99%까지 단축하였으며 (기하 평균 66.42%), 긴 시퀀스 (Longer sequences)에서 가장 큰 이득을 얻었습니다. 실제 데이터셋에서 RED-PIM은 정확도를 유지하거나 향상시키면서, 긴 문서의 경우 성능을 99.60%, 짧은 문서의 경우 13.44% 개선했습니다. 이러한 결과는 확장 가능하고 효율적인 transformer 추론을 위한 RED-PIM의 효과를 입증합니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.AR의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기