[논문 리뷰] TRON: Transformer Neural Network Acceleration with Non-Coherent Silicon Photonics

Salma Afifi, Febin Sunny외 2명|arXiv (Cornell University)|2023. 03. 22.

Neural Networks and Reservoir Computing인용 수 4

한 줄 요약

TRON은 BERT 및 비전 트랜스포머와 같은 트랜스포머 모델을 위한 첫 번째 실리콘 포토닉 신경망 가속기로, 비공진 실리콘 포토닉스를 활용하여 최신 전자 가속기 대비 최소 14배 높은 처리량과 8배 높은 에너지 효율성을 달성한다. 이는 자연어 처리 및 컴퓨터 비전 워크로드의 고성능 추론을 가능하게 한다.

ABSTRACT

Transformer neural networks are rapidly being integrated into state-of-the-art solutions for natural language processing (NLP) and computer vision. However, the complex structure of these models creates challenges for accelerating their execution on conventional electronic platforms. We propose the first silicon photonic hardware neural network accelerator called TRON for transformer-based models such as BERT, and Vision Transformers. Our analysis demonstrates that TRON exhibits at least 14x better throughput and 8x better energy efficiency, in comparison to state-of-the-art transformer accelerators.

연구 동기 및 목표

• 자연어 처리 및 컴퓨터 비전 분야에서 트랜스포머 기반 모델의 증가하는 계산 요구 사항을 해결한다.
• 주의 메커니즘의 고정밀도, 고대역폭 계산을 처리하는 데 있어 전자 가속기의 한계를 극복한다.
• 트랜스포머의 고유한 계산 패턴에 최적화된 포토닉 하드웨어 가속기 아키텍처를 설계한다.
• 최신 전자 가속기 대비 처리량과 에너지 효율성에서 뚜렷한 향상을 이룬다.
• 대규모 AI 추론에 적합한 비공진 실리콘 포토닉스의 실현 가능성과 성능 향상을 입증한다.

제안 방법

• 자기주의 메커니즘의 핵심이 되는 행렬 곱셈과 소프트맥스 연산에 최적화된 포토닉 가속기 아키텍처를 설계한다.
• 공진 시스템에 비해 시스템 복잡성과 비용을 줄이기 위해 비공진 실리콘 포토닉스를 활용한다.
• 제어 논리와 데이터 플로우 조율를 처리하기 위해 하이브리드 전자-포토닉 처리 파이프라인을 구현한다.
• 다중 주의 헤드와 계산 경로를 병렬화하기 위해 파장 분할 Multiplexing을 활용한다.
• 처리 요소 간 고대역폭, 저지연 통신을 위해 포토닉 인터커넥트를 통합한다.
• 변환 손실과 지연을 최소화하기 위해 주의 계산을 포토닉 도메인 연산에 최적화된 방식으로 매핑한다.

실험 결과

연구 질문

• RQ1 비공진 실리콘 포토닉스가 트랜스포머 신경망의 효율적 가속을 가능하게 할 수 있는가?
• RQ2 주의 기반 모델에 대해 포토닉 가속기와 전자 가속기 간 처리량과 에너지 효율성에서 어떤 차이가 있는가?
• RQ3 실세계 트랜스포머에 대해 포토닉 추론을 실현 가능하게 하기 위해 필요한 아키텍처적 및 신호 처리 혁신은 무엇인가?
• RQ4 포토닉 하드웨어가 자기주의 계산의 지연과 에너지 소비를 어느 정도 줄일 수 있는가?
• RQ5 확장 가능한 AI 가속기에서 비공진 광학을 사용할 경우 성능 상의 Trade-off는 무엇인가?

주요 결과

• TRON은 최신 전자 트랜스포머 가속기 대비 최소 14배 높은 처리량을 달성한다.
• 제안된 가속기는 기존 전자 솔루션 대비 8배 높은 에너지 효율성을 보여준다.
• 비공진 실리콘 포토닉스는 복잡한 위상 안정화가 필요 없이도 확장 가능하고 비용 효율적인 포토닉 가속기 설계를 가능하게 한다.
• 아키텍처는 BERT 및 비전 트랜스포머의 주의 메커니즘에서 요구하는 고정밀도, 고대역폭 요구 사항을 효과적으로 처리한다.
• 성능 향상은 포토닉 도메인 계산에 내재된 병렬성과 저지연 인터커넥트 덕분이다.
• 결과는 실리콘 포토닉스가 전통적인 전자 한계를 초월해 현대 딥 러닝 워크로드를 가속화하는 데 실현 가능하다는 것을 입증한다.