[논문 리뷰] Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations

Tri Dao, Albert Gu외 3명|arXiv (Cornell University)|2019. 03. 14.

Tensor decomposition and applications참고 문헌 48인용 수 31

한 줄 요약

이 논문은 구조적 선형 변환—예를 들어 DFT, DCT, 히터드, 그리고 컨볼루션—에 대해 $O(N\log N)$ 속도로 작동하는 자동으로 학습되는 빠른 알고리즘을 학습할 수 있는 미분 가능하고 매개변수화된 버터플라이 분해를 제안한다. 이는 희박하고 구조화된 블록 대각 행렬의 곱으로 표현된다. 이 방법은 엔드 투 엔드 머신 러닝 모델에서 최신 기술 성능을 달성하며, CIFAR-10에서 완전히 연결된 층보다 3.9점 높은 정확도를 기록하면서도 파rameter 수는 40배 줄이고 추론 속도는 4배 빠르게 한다.

ABSTRACT

Fast linear transforms are ubiquitous in machine learning, including the discrete Fourier transform, discrete cosine transform, and other structured transformations such as convolutions. All of these transforms can be represented by dense matrix-vector multiplication, yet each has a specialized and highly efficient (subquadratic) algorithm. We ask to what extent hand-crafting these algorithms and implementations is necessary, what structural priors they encode, and how much knowledge is required to automatically learn a fast algorithm for a provided structured transform. Motivated by a characterization of fast matrix-vector multiplication as products of sparse matrices, we introduce a parameterization of divide-and-conquer methods that is capable of representing a large class of transforms. This generic formulation can automatically learn an efficient algorithm for many important transforms; for example, it recovers the $O(N \log N)$ Cooley-Tukey FFT algorithm to machine precision, for dimensions $N$ up to $1024$. Furthermore, our method can be incorporated as a lightweight replacement of generic matrices in machine learning pipelines to learn efficient and compressible transformations. On a standard task of compressing a single hidden-layer network, our method exceeds the classification accuracy of unconstrained matrices on CIFAR-10 by 3.9 points -- the first time a structured approach has done so -- with 4X faster inference speed and 40X fewer parameters.

연구 동기 및 목표

• 수작업으로 설계된 방법 없이도 구조적 선형 변환에 대해 효율적이고 이차 이하의 알고리즘을 자동으로 학습하는 것.
• 플랫폼에 특화된 수작업 최적화 라이브러리(예: FFTW, cuFFT)에 대한 의존도를 줄이기 위해 빠른 변환을 엔드 투 엔드로 미분 가능한 학습이 가능하도록 하는 것.
• 구조화된 매개변수화가 실제 머신 러닝 벤치마크에서 비구조적 밀집 층을 능가할 수 있는지 확인하는 것.
• 희박한 행렬 분해를 통해 빠른 알고리즘을 학습하기 위해 필요한 최소한의 인덕티브 바이어스를 규명하는 것.
• 딥 러닝 파이프라인에 학습 가능한 빠른 변환 레이어를 통합하여 효율성과 성능을 향상시키는 것.

제안 방법

• 선형 변환을 $O(\log N)$개의 버터플라이 행렬 곱으로 매개변수화—구조화되고 희박한 블록 대각 행렬로, 총 파라미터 수는 $O(N)$이다.
• 목표 변환의 입력-출력 쌍에서 버터플라이 매개변수를 학습하기 위해 미분 가능한 최적화 프레임워크를 사용한다.
• 버터플라이 분해의 재귀적 분할 정복 구조를 활용하여 $O(N\log N)$ 행렬-벡터 곱셈을 가능하게 한다.
• 변환을 $O(\log N)$단계의 조합으로 표현하며, 각 단계는 치환과 크기가 $N/2^k$인 $2^k$개의 블록으로 구성된 블록 대각 행렬을 적용한다.
• 신경망에 버터플라이 레이어를 통합하여 완전히 연결된 층의 경량이고 압축 가능한 대체 수단으로 사용한다.
• 빠른 곱셈 알고리즘을 5줄의 파이썬 코드로 구현하여 어떤 하드웨어에서도 변환 전용 튜닝 없이도 효율적인 추론을 가능하게 한다.

실험 결과

연구 질문

• RQ1 일반적인, 미분 가능한 매개변수화가 DFT나 DCT와 같은 다양한 구조적 변환에 대해 $O(N\log N)$ 속도의 빠른 알고리즘을 학습할 수 있는가?
• RQ2 구조화된, 학습 가능한 행렬 분해가 정확도, 파라미터 효율성, 추론 속도 측면에서 비구조적 밀집 층을 얼마나 뛰어나게 성능을 높일 수 있는가?
• RQ3 동일한 버터플라이 매개변수화가 실제 크기까지 $N=1024$까지의 정확도로 알려진 빠른 알고리즘(예: 쿠리-트케이 FFT)을 기계 정밀도 수준에서 복원할 수 있는가?
• RQ4 학습된 버터플라이 변환의 성능이 매우 최적화된 수작업 튜닝된 커널들(예: cuFFT, cuDNN)과 비교해 어떻게 되는가?
• RQ5 버터플라이 레이어가 추가적인 파라미터를 크게 증가시키지 않으면서도 표준 딥 러닝 아키텍처인 ResNet18의 성능을 향상시킬 수 있는가?

주요 결과

• 이 방법은 $N$이 1024까지일 때 기계 정밀도 수준에서 $O(N\log N)$ 쿠리-트케이 FFT 알고리즘을 정확히 복원하며, 희박성 및 낮은 랭크 기반 베이스라인을 모두 능가한다.
• CIFAR-10에서 버터플라이 매개변수화를 사용한 단일 레이어 네트워크는 93.89%의 정확도를 기록—비구조적 완전히 연결된 층보다 3.9점 높은 성능—동시에 파라미터 수는 40배 줄였다.
• 버터플라이 레이어를 사용한 ResNet18은 CIFAR-10에서 정확도를 0.43점 향상시켰으며, 모델 파라미터 수는 0.07%만 증가했다.
• 학습된 버터플라이 행렬의 추론 속도는 밀집 GEMV보다 최대 2개의 지수 차이로 빠르며, CPU에서 전용 FFT/DCT 커널과 비교해 3~5배 이내였다.
• GPU에서 버터플라이 레이어를 사용한 학습은 밀집 GEMM보다 15% 더 빠르며, FFT 성능에 비해 40% 이내였으며, 경쟁 가능한 학습 효율성을 보였다.
• 이 방법은 컨볼루션 및 기타 구조적 변환을 학습하는 데에도 일반화되며, 빠른 곱셈 알고리즘에 변환 전용 최적화가 필요하지 않다.