[논문 리뷰] EcoFormer: Energy-Saving Attention with Linear Complexity

Jing Liu, Zizheng Pan외 3명|arXiv (Cornell University)|2022. 09. 19.

Advanced Neural Network Applications인용 수 16

한 줄 요약

EcoFormer는 쿼리와 키를 이산형 해싱을 통해 압축된 이진 코드로 매핑함으로써 쌍별 유사성을 유지하면서도 고에너지 소비 연산을 저에너지 누적 연산으로 대체하는 선형 복잡도이자 에너지 효율적인 어텐션 메커니즘을 제안한다. 이는 표준 어텐션 대비 73% 낮은 片상 에너지 소비를 달성했으며, ImageNet-1K에서 정확도 저하율은 0.33%에 불과하다.

ABSTRACT

Transformer is a transformative framework that models sequential data and has achieved remarkable performance on a wide range of tasks, but with high computational and energy cost. To improve its efficiency, a popular choice is to compress the models via binarization which constrains the floating-point values into binary ones to save resource consumption owing to cheap bitwise operations significantly. However, existing binarization methods only aim at minimizing the information loss for the input distribution statistically, while ignoring the pairwise similarity modeling at the core of the attention. To this end, we propose a new binarization paradigm customized to high-dimensional softmax attention via kernelized hashing, called EcoFormer, to map the original queries and keys into low-dimensional binary codes in Hamming space. The kernelized hash functions are learned to match the ground-truth similarity relations extracted from the attention map in a self-supervised way. Based on the equivalence between the inner product of binary codes and the Hamming distance as well as the associative property of matrix multiplication, we can approximate the attention in linear complexity by expressing it as a dot-product of binary codes. Moreover, the compact binary representations of queries and keys enable us to replace most of the expensive multiply-accumulate operations in attention with simple accumulations to save considerable on-chip energy footprint on edge devices. Extensive experiments on both vision and language tasks show that EcoFormer consistently achieves comparable performance with standard attentions while consuming much fewer resources. For example, based on PVTv2-B0 and ImageNet-1K, Ecoformer achieves a 73% on-chip energy footprint reduction with only a 0.33% performance drop compared to the standard attention. Code is available at https://github.com/ziplab/EcoFormer.

연구 동기 및 목표

• 특히 엣지 디바이스에서 높은 에너지 소비를 보이는 표준 트랜스포머의 문제를 해결한다.
• 기존의 이진화 방법이 어텐션 메커니즘에서 쌍별 유사성을 忽略하는 한계를 극복한다.
• 토큰 수준의 유사성 관계를 유지할 수 있도록 소프트맥스 어텐션 메커니즘에 특화된 이진화 전략을 개발한다.
• 성능 저하를 최소화하면서 어텐션 계산의 복잡도를 선형으로 유지한다.
• 에너지 효율적인 하드웨어 친화적 연산을 통해 자원 제약이 있는 엣지 디바이스에서 트랜스포머 모델의 확장 가능한 구현을 가능하게 한다.

제안 방법

• RBF 커널을 사용한 커널화된 해싱을 통해 고차원 쿼리와 키를 해밍 공간 내의 저차원 이진 코드로 매핑한다.
• 자기 지율적 방식으로 해시 함수를 학습하여 해밍 유사성(이진 코드의 내적)이 실제 어텐션 점수와 일치하도록 한다.
• 행렬 곱셈의 결합 법칙을 활용하여 어텐션을 이진 코드의 선형 복잡도 내적 연산으로 표현한다.
• 대부분의 곱셈-누적 연산을 단순한 누적 연산으로 대체하여 片상 에너지 소비를 줄인다.
• 표준 어텐션의 유사성 행렬을 해시 함수 학습을 위한 지도 신호로 사용한다.
• 학습 중 일반화 성능 향상을 위해 초수치 $ m $을 사용하는 서포트 샘플 전략을 도입한다.

실험 결과

연구 질문

• RQ1 자기 어텐션에서 쿼리와 키의 이진화가 통계적 양자화 오차 최소화를 넘어서 쌍별 유사성 관계를 유지하도록 맞춤형으로 설계될 수 있는가?
• RQ2 RBF 커널을 사용한 커널화된 해싱이 경쟁력 있는 성능 유지를 하면서도 선형 복잡도 어텐션을 가능하게 할 수 있는가?
• RQ3 FP32/FP16 곱셈-누적 연산을 이진 누적 연산으로 대체함으로써 트랜스포머의 에너지 소비를 얼마나 줄일 수 있는가?
• RQ4 서포트 샘플 수 $ m $의 변화가 제안된 방법의 성능 및 에너지 효율성에 어떤 영향을 미치는가?
• RQ5 EcoFormer는 시각 및 언어 벤치마크에서 표준 어텐션과 유사한 정확도를 확보하면서도 계산 및 에너지 비용을 크게 줄일 수 있는가?

주요 결과

• EcoFormer는 PVTv2-B0에서 표준 어텐션 대비 73% 낮은 片상 에너지 소비를 기록했으며, ImageNet-1K에서 정확도 저하율은 0.33%에 불과하다.
• CIFAR-100에서 $ m=25 $인 EcoFormer는 더 적은 서포트 샘플을 사용한 모델들을 능가하는 71.23%의 Top-1 정확도를 달성했다.
• 계산 복잡도를 $ \tfrac{1}{2}N^2 $ 에서 $ \tfrac{1}{2}N $으로 감소시켜 선형 시간 어텐션 계산을 가능하게 하였다.
• 다양한 벤치마크에서 다른 효율적인 어텐션 방법들과 비교해 EcoFormer는 유의미하게 낮은 지연 시간과 片상 에너지 소비를 보였다.
• ImageNet-1K에서의 초기 학습 결과로 EcoFormer는 PVTv2-B0에서 곱셈 연산을 72% 줄이고 에너지 소비를 72% 낮추면서도 높은 성능를 유지하였다.
• 이진화로 인한 성능 저하는 비가역적인 부호 함수에서의 기울기 근사 오차에 기인하며, 이는 고도의 최적화 기법으로 완화될 수 있다.