[논문 리뷰] ABC: Attention with Bounded-memory Control

Hao Peng, Jungo Kasai외 6명|arXiv (Cornell University)|2021. 10. 06.

Topic Modeling인용 수 7

한 줄 요약

이 논문은 선형 복잡도를 달성하기 위해 메모리 크기를 제한함으로써 효율적인 어텐션 메커니즘을 일반화하는 통합 프레임워크인 ABC(Attention with Bounded-Memory Control)를 소개한다. ABC$_{\text{MLP}}$는 학습된, 문맥 기반의 메모리 제어 메커니즘으로, 특히 장문의 시퀀스에서 기존 방법보다 높은 정확도와 효율성을 보이며, 표준 트랜스포머에 비해 정확도 손실가 최소화된다.

ABSTRACT

Transformer architectures have achieved state-of-the-art results on a variety of sequence modeling tasks. However, their attention mechanism comes with a quadratic complexity in sequence lengths, making the computational overhead prohibitive, especially for long sequences. Attention context can be seen as a random-access memory with each token taking a slot. Under this perspective, the memory size grows linearly with the sequence length, and so does the overhead of reading from it. One way to improve the efficiency is to bound the memory size. We show that disparate approaches can be subsumed into one abstraction, attention with bounded-memory control (ABC), and they vary in their organization of the memory. ABC reveals new, unexplored possibilities. First, it connects several efficient attention variants that would otherwise seem apart. Second, this abstraction gives new insights--an established approach (Wang et al., 2020b) previously thought to be not applicable in causal attention, actually is. Last, we present a new instance of ABC, which draws inspiration from existing ABC approaches, but replaces their heuristic memory-organizing functions with a learned, contextualized one. Our experiments on language modeling, machine translation, and masked language model finetuning show that our approach outperforms previous efficient attention models; compared to the strong transformer baselines, it significantly improves the inference time and space efficiency with no or negligible accuracy loss.

연구 동기 및 목표

• 다양한 효율적인 어텐션 메커니즘을 하나의 추상화로 통합하여, 서로 다른 것처럼 보이는 접근 방식들 사이의 숨겨진 구조적 원리를 드러내는 것.
• 장문의 시퀀스에서 표준 트랜스포머의 제곱형 계산 비용 문제를 해결하기 위해 메모리 크기를 제한하면서도 성능을 유지하는 것.
• 히우레스틱 기반 방법보다 효율성과 정확도의 균형을 향상시키기 위해 학습 가능한, 문맥 인식 메모리 제어 전략을 개발하는 것.
• 이전에는 인접한 어텐션에 적용 불가능하다고 여겨졌던 방법들, 예를 들어 Linformer가 ABC 프레임워크를 통해 어떻게 적응 가능한지 보여주는 것.
• 표준 자기어텐션의 실용적, 즉각적인 대체 수단을 제공하여 정확도 손실 없이 유의미한 속도 향상과 메모리 절감을 달성하는 것.

제안 방법

• ABC는 고정 크기의 메모리가 키-밸류 쌍을 저장하는 제한된 메모리 시스템으로 어텐션을 프레임워크화하며, 쿼리 접근은 메모리 구성 전략에 의해 제어된다.
• 이 프레임워크는 메모리 구성 방식을 다른 제어 메커니즘으로 간주함으로써 기존 방법들—예: 희소 어텐션, 국소 어텐션, Linformer—을 추상화한다.
• ABC$_{\text{MLP}}$는 문맥 기반 입력 특징에 따라 어떤 토큰을 메모리에 저장할지, 어디에 저장할지를 결정하는 학습 가능한 신경망 기반 컨트롤러를 도입한다.
• 컨트롤러는 다층 퍼셉트론(Multi-Layer Perceptron, MLP)을 사용하여 메모리 쓰기 우선순위를 계산함으로써 동적이고 문맥 민감한 메모리 관리가 가능하다.
• ABC는 인과적 및 비인과적 어텐션 모두를 지원하여 선형 복잡도를 유지하면서도 효율적인 자동회귀 디코딩과 시퀀스 인코딩을 가능하게 한다.
• 어텐션 계산은 제한된 메모리에 대해 소프트맥스 정규화된 읽기 연산을 사용하여 일정 시간 및 일정 공간의 액세스 오버헤드를 보장한다.

실험 결과

연구 질문

• RQ1 다양한 효율적인 어텐션 메커니즘은 하나의 추상화로 통합될 수 있는가? 이는 그들의 공통된 구조적 원리를 드러내는가?
• RQ2 학습 가능한, 문맥 인식 메모리 제어 전략은 히우레스틱 기반 메모리 구성 방식보다 정확도와 효율성 측면에서 뛰어나게 되는가?
• RQ3 이전에 인과 어텐션에 적용 불가능하다고 여겨졌던 기존 방법들, 예를 들어 Linformer는 ABC 프레임워크를 통해 어떻게 호환 가능한가?
• RQ4 메모리 크기가 다양한 어텐션 유형과 작업에서 모델 성능에 미치는 영향은 어떠한가?
• RQ5 ABC$_{\text{MLP}}$는 추론 시간과 메모리 사용량을 줄이면서 표준 트랜스포머 성능을 따라하거나 초월할 수 있는가?

주요 결과

• ABC$_{\text{MLP}}$는 RFA보다 최대 1.25배 빠른 디코딩 속도를 기록했으며, T2R보다 20% 더 빠르며, 메모리 오버헤드가 크게 감소했다.
• 8-8 메모리 크기(교차 및 인과적)를 사용한 ABC$_{\text{MLP}}$는 최고 성능 설정에 비해 BLEU 점수로 1.1점 뒤져도 높은 효율-정확도 균형을 보였다.
• 텍스트 인코딩에서 ABC$_{\text{MLP}}$는 표준 트랜스포머 베이스 모델 대비 메모리 사용량을 50% 감소시키고 추론 속도를 1.5배 향상시켰다.
• 이전에는 인과 어텐션에 적용 불가능하다고 여겨졌던 Linformer는 ABC 추상화를 통해 호환 가능하다는 것이 입증되었으며 기계 번역 작업에서 경쟁적인 성능을 보였다.
• ABC$_{\text{MLP}}$는 더 작은 메모리 크기로도 다른 모든 ABC 변종보다 정확도에서 뛰어나며, 문맥 기반 메모리 제어의 유용성을 확인했다.
• ABC$_{\text{MLP}}$의 효율성 향상은 특히 장문의 시퀀스에서 두드러지게 나타나, 장문 시퀀스 응용 분야에서 강력한 渐近적 우수성을 보였다.