https://arxiv.org/pdf/2406.14909
간단 요약
"MOA: Mixture of Sparse Attention for Automatic Large Language Model Compression" 논문은 대규모 언어 모델(LLM)의 메모리 및 연산 효율성을 향상시키기 위한 Mixture of Attention (MoA) 기법을 제안합니다. MoA는 각 레이어와 어텐션 헤드에 대해 서로 다른 스파스(sparse) 패턴을 적용하여 최적의 압축 방안을 자동으로 찾아내는 것을 목표로 합니다.
이 논문은 다중 모델 자체를 사용하는 것이 아니라, 하나의 모델 내에서 서로 다른 어텐션 패턴을 활용하여 효율성을 최적화하는 방식으로, 여러 모델을 다루는 MoA와는 다른 접근 방식을 취하고 있습니다.
MOA: Mixture of Sparse Attention for Automatic Large Language Model Compression 논문에서 제안하는 MoA(Mixture of Attention) 기법은 대규모 언어 모델(LLM)의 효율적인 압축을 위해 각 레이어와 어텐션 헤드에 개별화된 스파스 패턴을 적용하는 방법론입니다. 아래는 MoA의 주요 방법론을 단계별로 자세히 설명한 내용입니다.
MoA의 압축 절차는 다음과 같은 단계로 구성됩니다:
MoA는 각 어텐션 헤드와 레이어에 대해 개별화된 스파스 패턴을 적용하는 접근 방식을 통해, 메모리 절약 및 연산 효율을 크게 개선하고 대규모 언어 모델이 긴 문맥을 처리할 때 발생하는 문제를 해결합니다.
스파스 어텐션은 긴 문맥을 다루는 대규모 언어 모델(LLM)의 메모리와 처리 성능 요구 사항을 효과적으로 줄일 수 있습니다. 기존 방법은 일반적으로 동일한 스파스 패턴을 모든 어텐션 헤드와 입력 길이에 동일하게 적용하는 균일한 스파스 어텐션 마스크를 사용합니다. 하지만, 이러한 균일한 접근법은 LLM이 내재한 다양한 어텐션 패턴을 포착하지 못해 정확성과 처리 성능 간의 최적화를 어렵게 만듭니다. 이를 해결하기 위해, 우리는 Mixture of Attention (MoA)을 제안하여, 각 어텐션 헤드와 레이어에 대해 다른 스파스 어텐션 구성을 자동으로 최적화합니다. MoA는 다양한 어텐션 패턴과 이를 입력 시퀀스 길이에 맞추는 규칙을 포함한 탐색 공간을 구축하고, 프로파일링을 통해 최적의 스파스 어텐션 압축 계획을 찾아냅니다. MoA는 입력 크기에 맞춰 조정되며, 일부 어텐션 헤드는 긴 시퀀스를 확장하여 처리하고, 다른 헤드는 고정된 짧은 문맥에 집중하는 방식으로 작동합니다.
실험 결과, MoA는 Vicuna-{7B, 13B} 및 Llama3-{8B, 70B} 모델에서 동일한 평균 어텐션 범위를 사용해 유효한 문맥 길이를 3.9배 늘리고, 검색 정확도를 균일한 어텐션 기반보다 1.5-7.1배 높였습니다. 또한, MoA는 스파스 모델과 밀집 모델 간의 성능 차이를 줄여주며, 긴 문맥 이해 벤치마크에서 성능 저하를 5% 이내로 유지했습니다. MoA는 GPU 메모리 사용량을 1.2-1.4배 감소시키고, FlashAttention2 및 vLLM 대비 6.6-8.2배와 1.7-1.9배의 디코딩 처리 성능 향상을 달성하며 성능에 미치는 영향을 최소화합니다. 코드와 추가 자료는 GitHub에서 확인할 수 있습니다.
대규모 언어 모델(LLM)은 다양한 응용 분야에서 놀라운 성능을 발휘하고 있으며, 그 중심에는 토큰 간 상호작용을 계산하여 문맥을 이해하게 하는 어텐션 메커니즘이 있습니다. 입력 길이를 확장하는 것은 LLM의 능력을 강화하는 데 필수적인 요소로, 정보 검색, 요약, Few-shot 학습, 질문 응답 등 다양한 과제에서 중요한 역할을 합니다. 그러나, 이러한 어텐션 계산과 Key-Value 캐시(KV-Cache)는 점점 더 큰 효율성 문제를 초래합니다.
기존 연구들은 생성형 LLM에서 긴 문맥의 효율성 문제를 해결하기 위해 스파스 어텐션 방법을 제안했습니다. 이 방법들은 보통 모든 어텐션 헤드와 입력 길이에 동일한 고정된 범위의 슬라이딩 윈도우 마스크를 사용하여 지역 문맥에만 주의를 기울이게 합니다. 이러한 방식은 긴 입력을 처리할 수 있게 하면서도 계산량과 KV 캐시 오버헤드를 제한하지만, 모든 어텐션 헤드가 동일한 스파스 마스크를 사용하는 것은 각 헤드가 필요로 하는 다양한 범위를 반영하지 못합니다. 특히, 일부 어텐션 헤드는 전체 시퀀스를 처리하기 위해 더 넓은 범위가 필요할 수 있으며, 다른 어텐션 헤드는 짧고 지역적인 문맥에만 주의를 기울여도 충분할 수 있습니다.
MoA는 이러한 문제를 해결하기 위해 이질적 엘라스틱 규칙(Heterogeneous Elastic Rules)을 통해 각 어텐션 헤드마다 다른 스파스 패턴을 적용하는 자동화된 압축 방법을 제안합니다. MoA는 각 어텐션 위치가 예측 손실에 미치는 영향을 프로파일링하기 위해 Gradient-based Profiling을 사용하며, 이를 통해 각 모델 레이어와 어텐션 헤드에 적합한 스파스 어텐션 구성을 적용합니다. MoA는 긴 범위의 의존성을 가진 데이터셋과 원본 밀집 모델의 출력을 참조하여, 정확한 프로파일링을 통해 더 나은 압축 결과를 얻을 수 있도록 합니다.
MoA의 주요 기여점:
실험 결과, MoA는 FlashAttention2 대비 6.6-8.2배, vLLM 대비 1.7-1.9배의 처리 성능을 보여주며, 검색 정확도의 손실은 평균 1% 이하에 그쳤습니다.
Multi-Head Self Attention (MHA) 메커니즘은 대규모 언어 모델(LLM) 기능에 매우 중요합니다. 이 메커니즘은 입력 시퀀스를 선형 변환을 통해 쿼리(Q), 키(K), 값(V) 행렬로 변환하는 것으로 시작합니다. 이 행렬들은 이전 시퀀스의 캐시된 K와 V (KV-Cache)와 함께 어텐션 행렬(A)을 계산하는 데 사용됩니다. 이 계산은 인과 마스크(M)에 의해 수정되어 자가회귀 특성을 유지하며, 아래와 같은 식으로 출력(O)이 계산됩니다:
LLM의 자가회귀적 추론은 두 단계로 이루어집니다: Prefill 단계와 Decode 단계입니다. Prefill 단계에서는 모델이 전체 입력 시퀀스를 처리하여 첫 번째 응답 토큰을 생성합니다. 이후 Decode 단계에서는 새로 생성된 토큰과 이전에 캐시된 K와 V 행렬을 사용해 다음 토큰을 생성하고, 생성이 완료될 때까지 이 과정을 반복합니다. 이 반복 과정은 메모리와 계산 요구량을 증가시키며, 특히 KV-Cache가 커지면서 이러한 부담이 커집니다.
어텐션 관련 계산과 메모리 비용을 줄이기 위해 다양한 효율적인 방법들이 제안되었습니다. 한 가지 방법은 동적 스파스 어텐션 마스크를 사용해, Prefill 단계에서 어텐션 계산을 선택적으로 건너뛰는 것입니다. 예를 들어, Pagliardini et al. (2023), Qu et al. (2022), Roy et al. (2021), Wang et al. (2021) 등의 연구는 입력 시퀀스에 따라 계산량을 줄이는 동적 스파스 어텐션을 제안합니다. 하지만, 동적 프리필은 복잡한 제어와 계산 흐름이 필요해 특정 하드웨어에서만 벽시계 시간을 단축할 수 있습니다.
또한 Decode 단계에서 KV-Cache를 동적으로 줄이는 방법도 제안되었습니다. 예를 들어, Anagnostidis et al. (2023)과 Zhang et al. (2023) 등은 추가적인 어텐션 스코어 계산이나 메모리 스왑을 통해, 이전에 저장된 KV-Cache를 최적화하는 방식을 제시했습니다. 하지만 이러한 동적 KV-Cache 축소는 광범위한 재학습이 필요할 수 있습니다.
정적 스파스 어텐션은 모든 처리 문장에 대해 미리 정의된 마스크를 일관되게 적용하는 방식입니다. 정적 스파스 어텐션은 계산 흐름이 고정되어 있어 일반적으로 GPU 효율성이 높고, BERT와 같은 언어 이해 모델에서는 다양한 마스크를 사용합니다. 그러나 생성형 LLM에서는 고정-범위 슬라이딩 윈도우 마스크가 일반적으로 사용되며, 초기 몇 개의 토큰에 대해서만 전역 어텐션을 허용합니다. 이러한 지역 어텐션 패턴을 사용하면 현재 어텐션 범위 밖의 KV-Cache를 제거할 수 있어 긴 시퀀스의 메모리를 절약할 수 있습니다.
하지만, 다양한 어텐션 헤드와 입력 길이에 대해 일관된 정적 마스크를 사용하는 것은 모델 및 데이터 특성을 반영하지 못하여 긴 시퀀스에서의 최적 성능을 보장하지 못할 수 있습니다. 우리의 방법론은 이와 유사한 방식이지만, 성능 한계를 극복하면서도 효율성과 학습이 필요 없다는 장점을 유지합니다.
스파스 어텐션 외에도 전통적인 어텐션을 대체하기 위한 새로운 메커니즘들이 제안되었습니다. 예를 들어, Gu & Dao (2023), Peng et al. (2023) 등의 연구들은 긴 시퀀스 모델링을 위해 새로운 구조를 제시하지만, 이러한 방법들은 일반적인 트랜스포머와는 다른 가중치를 필요로 하여 LLM에 재훈련 부담을 증가시킵니다.
기존 연구에서는 LLM의 가속화를 위한 다양한 프레임워크와 커널 수준의 최적화도 제안되었습니다. 이와 같은 시스템 최적화는 우리의 연구와 상호 보완적이며, 추가적인 효율성 향상을 위해 통합할 수 있습니다.

우선, 사전 학습된 LLM에서 어텐션 패턴의 이질성을 설명합니다. 이를 바탕으로 Mixture-of-Attention (MoA) 방법을 위한 탐색 공간을 정의합니다.
이질적인 어텐션 패턴: LLM의 다양한 어텐션 헤드들은 각기 다른 어텐션 패턴을 보입니다. 예를 들어, 첫 번째 헤드는 좁은 범위의 슬라이딩 윈도우로 지역적 문맥에 초점을 맞추는 반면, 세 번째 헤드는 거의 전체 입력을 포괄하여 전역적 어텐션을 제공합니다. 이러한 어텐션 범위는 대부분 다양한 작업과 데이터셋에 걸쳐 일정하게 유지됩니다. 예를 들어, 동일한 슬라이딩 윈도우 스파스 어텐션 마스크를 적용하면 검색 정확도에서 65%의 변동이 발생합니다. 이처럼 각 어텐션 헤드가 다양한 정보를 포착하도록 설계되어 있으며, 일부 연구에 따르면 특정 어텐션 헤드는 전역 텍스트 검색에 특화되어 있습니다.
이질적 탄력 규칙: 특정 길이에서의 이질성 외에도, 입력 길이에 따라 다른 탄력적 반응을 보입니다. 예를 들어, 짧은 입력에서는 두 번째와 세 번째 헤드가 전역 어텐션을 보이지만, 입력 길이가 증가하면 두 번째 헤드는 중간 범위의 지역적 초점으로 유지되고, 세 번째 헤드는 계속해서 전역적 범위를 확장합니다. 이러한 발견을 활용하여 MoA는 이질적 탄력 규칙을 탐색 공간에 포함합니다.
MoA 마스크를 설계할 때 LLM 어텐션 패턴의 이질적이고 탄력적인 특성을 고려합니다. 하드웨어 친화적인 슬라이딩 윈도우 마스크를 기본 스파스 어텐션 마스크로 채택하며, 초기 몇 개의 토큰(예: MoA에서 64개 토큰)은 마스킹되지 않습니다. 어텐션 범위 S_h 는 다음과 같이 정의됩니다:
여기서 와 는 특정 어텐션 헤드의 기본 범위와 입력 길이에 따른 확장 비율을 제어하는 하이퍼파라미터입니다. 각 어텐션 헤드는 여러 개의 및 옵션에서 선택될 수 있으며, 예를 들어 7B 모델의 경우 탐색 공간은 매우 커질 수 있습니다. 이를 위해 최적의 와 를 효율적으로 선택할 수 있는 자동 파이프라인을 설계했습니다.
MoA 자동 압축 파이프라인은 훈련된 LLM과 캘리브레이션 데이터셋을 사용하여 시작됩니다. MoA는 캘리브레이션 데이터셋의 다양한 입력 시퀀스에 대해 각 어텐션 값이 모델의 예측 손실에 미치는 영향을 프로파일링합니다. 각 후보 탄력 규칙에 따른 정확도 손실을 계산하여 최적의 규칙을 선택합니다.
MoA는 사전 학습된 LLM의 최종 예측 손실에 각 어텐션 값이 미치는 영향을 정량화합니다. 이때 주어진 캘리브레이션 데이터셋에 대해 어텐션 행렬 A 와 그에 대한 그라디언트 를 사용하여 어텐션 값이 마스킹될 경우 손실 변화를 추정합니다. 손실 변화 는 다음과 같이 1차 Taylor 전개를 사용하여 근사됩니다:
어텐션 값 가 마스킹될 때의 영향도 행렬 는 다음과 같이 정의됩니다:
각 어텐션 헤드의 평균 영향을 계산하여, 특정 입력 길이에 따라 각 후보 탄력 규칙이 적용될 때의 정확도 손실을 계산할 수 있습니다.
MoA는 여러 시퀀스 길이에 대해 정확도 손실을 최소화하면서 밀도 예산을 충족하는 최적의 탄력 규칙을 자동으로 선택합니다. 프로파일링 결과에 기반하여 다양한 길이에서 최적의 정확도 손실을 보이는 Pareto 최적 압축 계획을 식별하고, 보이지 않은 길이에 대한 일반화가 잘 되도록 최적 계획을 선택합니다.
이를 통해 MoA는 사용자가 정의한 밀도 제약을 준수하면서, 어텐션 스파스화로 인한 정확도 손실을 최소화합니다.
간략한 실험 설정을 설명하며, 자세한 사항은 Appendix A에 있습니다.
비교 대상: MoA는 최신 스파스 어텐션 방법들과 비교됩니다. 여기에는 StreamingLLM, InfLLM, H2O가 포함됩니다. MoA와 StreamingLLM에서는 KV-Cache 길이가 스파스 프리필 단계에서 어텐션 범위와 동일하며, 반면 H2O는 밀집된 프리필을 사용합니다. H2O와 InfLLM은 동적으로 KV-Cache를 결정하기 위한 추가 계산이 필요합니다.
모델과 벤치마크: Vicuna-7b, 13b 및 Llama-3-8b, 70b 모델에서 평가했습니다. 긴 문맥 검색을 위해 LongEval 데이터셋을 사용했고, 긴 문맥 이해 테스트에는 LV-Eval과 LongBench를 사용했습니다. 성능 실험은 A100-80GB GPU 8개에서 24시간 동안 수행되며, OOM (Out-Of-Memory) 및 OOT (Out-Of-Time) 조건이 기록됩니다.
MoA 설정: 각 모델 레이어에 대해 최대 두 개의 규칙만 사용하여 효율성을 유지합니다. MultiNews 데이터셋에서 2k, 4k, 8k 길이로 프로파일링한 후 12k에서 최적의 압축 계획을 선택합니다. 모델은 별도의 파인튜닝 없이 동일한 계획을 사용합니다.
MoA는 다양한 밀도에서 문맥 검색 정확도와 디코드 처리 속도에서 Pareto 효율을 개선합니다. MoA는 정적 어텐션 설계 덕분에 H2O와 InfLLM에 비해 1.6배에서 18.1배의 처리 속도 향상을 보입니다. 모든 밀도에서 MoA가 더 높은 검색 정확도를 달성하며, 이후 섹션에서는 MoA의 성능과 효율성을 더욱 폭넓게 평가합니다.
MoA는 다양한 모델 크기와 벤치마크에서 최신 스파스 어텐션 방법보다 우수한 성능을 보이며, 밀집 모델과 비교해도 50% 밀도에서 유사한 성능을 유지합니다.
긴 문맥 검색: Table 4에 따르면 MoA는 최대 8%의 상대적 정확도 감소를 나타내며, 이는 StreamingLLM, InfLLM, H2O가 각각 87%, 58%, 44%의 정확도 감소를 보이는 것과 대조적입니다. MoA는 평균적으로 1% 미만의 상대적 정확도 감소를 나타내어, 51%, 41%, 20%를 나타낸 다른 방법들에 비해 우수합니다.
긴 문맥 이해: Table 4에 따르면, MoA는 LV-Eval과 LongBench 벤치마크에서 상대 성능 저하를 각각 5%, 3%로 줄였으며, 이는 StreamingLLM의 36% 및 18%보다 훨씬 낮습니다. 다른 방법들에 비해 MoA는 일관된 성능을 유지하며, 다양한 작업에 대해 원래 밀집 모델에 필적하는 포괄적인 성능을 달성했습니다.
MoA는 실행 시간 효율성이 높으며, 일회성 압축 오버헤드가 관리 가능한 수준입니다.
실행 시간 효율성: Table 5에 따르면, MoA는 FlashAttention2 대비 6.6배에서 8.2배의 디코드 처리 속도 향상을 보이며, H2O와 InfLLM 대비 1.2배에서 4.0배의 속도 향상을 보여줍니다. 또한 MoA는 GPU 메모리 사용량을 1.2배에서 1.4배 줄입니다. 이 처리 속도 증가는 KV-Cache 크기의 고정, 어텐션 계산 감소, 배치 크기 증가, 그리고 CUDA 커널 구현 덕분입니다.
압축 파이프라인 효율성: Vicuna-7B 및 13B 모델의 경우, MoA는 2시간 내에 자동 압축 파이프라인을 완료하며, Llama3-70B 모델의 경우 8.5시간의 실시간과 34.7시간의 GPU 시간이 필요합니다.
MoA는 각 헤드에 대해 다양한 탄력적 규칙을 발견했습니다. 초기와 중간 레이어에서는 높은 밀도의 마스크가 필요하고, 마지막 레이어에서는 대부분 낮은 밀도를 필요로 합니다. 이는 동일한 레이어 내에서도 헤드마다 다른 밀도가 필요함을 나타냅니다.
MoA는 각 어텐션 헤드와 입력 길이에 대해 이질적 탄력 마스크를 자동으로 선택하여, LLM의 유효한 문맥 길이를 3.9배 확장하며, 50% 밀도에서 6배 이상의 처리 속도 증가와 밀집 모델과 유사한 성능을 달성했습니다.
제약 사항 및 향후 연구: 극히 낮은 밀도에서는 MoA의 성능이 저하될 수 있습니다. 이를 해결하기 위해 동적 MoA 방법을 설계하거나 비선형 탄력 규칙을 사용하는 것도 향후 연구 과제로 제시되었습니다.