[논문 리뷰] Simple Hardware-Efficient Long Convolutions for Sequence Modeling

Daniel Y. Fu, Elliot L. Epstein외 6명|arXiv (Cornell University)|2023. 02. 13.

Topic Modeling인용 수 6

한 줄 요약

이 논문은 시계열 모델링을 위한 장거리 컨볼루션 커널을 직접 파arameter화하는 방법을 제안하며, 단순한 정규화(예: 커널 가중치를 압축하는 것)를 통해 매끄러운 커널을 생성함으로써 상태공간모델(SSM) 수준의 성능을 달성함을 보여준다. FlashButterfly는 브로드캐스트 분해를 활용하는 I/O에 민감한 알고리즘으로, 장거리 컨볼루션을 cuFFT 대비 2.2배 빠르게 하고 64K 길이의 시퀀스에서 학습 속도를 7.2배 빠르게 하며, Path256 및 뇌 fMRI 데이터를 포함한 다양한 벤치마크에서 SSM과 트랜스포머를 능가한다.

ABSTRACT

State space models (SSMs) have high performance on long sequence modeling but require sophisticated initialization techniques and specialized implementations for high quality and runtime performance. We study whether a simple alternative can match SSMs in performance and efficiency: directly learning long convolutions over the sequence. We find that a key requirement to achieving high performance is keeping the convolution kernels smooth. We find that simple interventions--such as squashing the kernel weights--result in smooth kernels and recover SSM performance on a range of tasks including the long range arena, image classification, language modeling, and brain data modeling. Next, we develop FlashButterfly, an IO-aware algorithm to improve the runtime performance of long convolutions. FlashButterfly appeals to classic Butterfly decompositions of the convolution to reduce GPU memory IO and increase FLOP utilization. FlashButterfly speeds up convolutions by 2.2$ imes$, and allows us to train on Path256, a challenging task with sequence length 64K, where we set state-of-the-art by 29.1 points while training 7.2$ imes$ faster than prior work. Lastly, we introduce an extension to FlashButterfly that learns the coefficients of the Butterfly decomposition, increasing expressivity without increasing runtime. Using this extension, we outperform a Transformer on WikiText103 by 0.2 PPL with 30% fewer parameters.

연구 동기 및 목표

• 복잡한 초기화나 특수한 구현 없이도 직접 파arameter화된 장거리 컨볼루션 커널이 상태공간모델(SSM)의 성능을 재현할 수 있는지 조사하기 위해.
• 현대 하드웨어, 특히 GPU 메모리 I/O 병목 현상과 활용되지 않는 FLOP 용량으로 인한 장거리 컨볼루션의 런타임 비효율성 문제를 해결하기 위해.
• 브로드캐스트 분해와 같은 구조적 행렬 분해가 장거리 컨볼루션의 효율성과 표현력 향상에 기여할 수 있는지 탐색하기 위해.
• 언어, 시각, 신경영상과 같은 다양한 모odalities에 걸쳐 제안된 방법의 일반화 능력을 평가하기 위해.
• 단순한 정규화와 효율적인 알고리즘 설계를 통해 표준 컨볼루션과 SSM 간의 성능 격차를 메우고, 장거리 시퀀스에서의 빠른 학습을 가능하게 할 수 있는지 증명하기 위해.

제안 방법

• 주파수 도메인에서의 매끄러움과 흐릿함을 증진하기 위해 커널 가중치에 Squash 연산자를 적용하여 희박성과 매끄러움을 강제한다.
• 시간 도메인에서 Smooth 연산자를 사용하여 커널 가중치를 추가로 정규화하고 주파수 도메인의 매끄러움을 향상시킨다.
• 블록 희박한 브로드캐스트 행렬 분해를 사용하여 FFT 기반 컨볼루션을 재작성함으로써 메모리 I/O를 줄이고 FLOP 활용도를 향상시키는 I/O에 민감한 알고리즘인 FlashButterfly를 개발한다.
• 브로드캐스트 분해를 통해 컨볼루션을 희박한 행렬 곱셈의 시퀀스로 재구성함으로써 GPU 행렬 곱셈 유닛을 효율적으로 활용한다.
• 분해 계수를 학습 가능하게 하는 학습된 브로드캐스트 확장 기법을 도입하여 표현력을 높이되, 추론 비용은 증가시키지 않는다.
• 결과로 도출된 아키텍처를 S4 및 H3와 같은 모델의 SSM 레이어에 직접 통합 가능한 교체 가능한 대체품으로 활용하고, LRA, CIFAR, WikiText103, fMRI 등 다양한 벤치마크에서 평가한다.

실험 결과

연구 질문

• RQ1 복잡한 초기화나 특수한 커널 없이도 직접 파arameter화된 장거리 컨볼루션 커널이 장거리 시퀀스 모델링 작업에서 SSM의 성능을 재현할 수 있는가?
• RQ2 장거리 컨볼루션을 SSM과 경쟁 가능하게 만들기 위해 어떤 단순한 정규화 기법이 매끄러운 커널 행동을 촉진할 수 있는가?
• RQ3 I/O에 민감한 알고리즘인 FlashButterfly가 현대 GPU에서 장거리 컨볼루션의 런타임 효율성을 크게 향상시킬 수 있는가?
• RQ4 구조적 브로드캐스트 분해를 사용하면 장거리 컨볼루션에서 속도 향상과 더불어 표현력 향상이 가능해지는가?
• RQ5 fMRI 시퀀스 모델링과 같은 도전적인 실제 응용 분야에서 장거리 컨볼루션은 트랜스포머와 SSM을 능가할 수 있는가?

주요 결과

• 단순한 정규화—특히 커널 가중치를 압축하고 스무딩하는 것—은 주파수 도메인에서 매끄러운 커널을 생성하며, LRA에서 S4 수준의 성능을 회복한다. 랜덤 초기화 상태에서도 S4와 동일한 성능을 달성한다.
• 이미지 분류에서 장거리 컨볼루션은 순차적 CIFAR에서 S4보다 0.3% 높은 정확도를 기록했고, 2D CIFAR에서 S4ND-ISO에 비해 0.8% 이내의 정확도를 기록했다.
• 언어 모델링에서 장거리 컨볼루션은 PILE에서 H3 성능을 0.3 PPL 이내로 재현했으며, OpenWebText에서도 0.3 PPL 이내의 성능을 기록했다.
• 뇌 fMRI 모델링에서 장거리 컨볼루션은 트랜스포머와 SSM을 각각 0.14점, 0.16점의 MAE 향상으로 능가했다.
• FlashButterfly는 cuFFT 대비 장거리 컨볼루션의 처리 속도를 2.2배 빠르게 하고, Path256(64K 시퀀스 길이)에서의 학습을 가능하게 하여 기존 연구 대비 29.1점 높은 정확도와 7.2배 빠른 학습 속도를 달성하며 최신 기술 수준의 성능을 확보했다.
• 학습된 브로드캐스트 확장 기법은 런타임 비용을 증가시키지 않으면서도 표현력을 높여, WikiText103에서 트랜스포머보다 0.2 PPL 낮은 성능을 기록했고, 파aram터 수는 30% 적게 사용했다.