Published on

대규모 언어 모델 추론 효율성 향상: MFA의 KV 캐시 감소 전략

저자
  • avatar
    이름
    Ajax
    Twitter

대규모 언어 모델(LLM) 추론 효율성 문제

최근 대규모 언어 모델(LLM)의 사용이 증가하면서 효율적인 대규모 추론이 중요한 과제로 떠올랐습니다. 기존 어텐션 메커니즘 내의 Key-Value (KV) 캐시는 배치 크기와 시퀀스 길이에 따라 선형적으로 증가하여 메모리 사용량이 급증하고 LLM의 확장성을 저해하는 주요 원인입니다. MQA, GQA, MLA와 같은 대안들이 등장했지만, 엄격한 메모리 제약 조건에서 성능을 유지하거나 엔지니어링 복잡성을 야기하는 문제가 있었습니다.

Multi-matrix Factorization Attention (MFA)의 등장

Stepes, Tsinghua University 등의 연구진은 Multi-matrix Factorization Attention (MFA) 및 MFA-Key-Reuse (MFA-KR)라는 새로운 어텐션 메커니즘 아키텍처를 제시했습니다. 이 메커니즘은 언어 모델 추론 비용을 크게 줄이면서도 성능을 향상시킵니다. MFA와 MFA-KR은 MLA보다 뛰어난 성능을 보이며, 기존 MHA 성능과 유사하면서도 KV 캐시 사용량을 최대 93.7%까지 줄입니다. MFA는 단순하고 재현하기 쉬우며, 하이퍼파라미터에 대한 민감도가 낮고 다양한 Pos-embedding 방법과 호환됩니다.

MFA 접근 방식 및 분석

연구팀은 어텐션 메커니즘의 일반적인 설계와 용량을 분석하여 용량과 관련된 두 가지 중요한 차원을 확인했습니다. 이 분석을 통해 새로운 분석 방법과 설계 원칙을 개발했습니다. 그들은 다양한 MHA 변형을 이해하기 위한 통합 프레임워크로 Generalized Multi-Head Attention (GMHA) 개념을 도입했습니다. 또한, 추론 관점에서 키-값의 계산 및 저장 방식을 탐구하고 분해 관점에서 모델 용량을 검토했습니다. Fully Parameterized Bilinear Attention (FPBA)을 이론적 성능 상한으로 설정하고, MHA와 그 변형이 FPBA의 저랭크 분해임을 발견했습니다.

MQA 및 MLA와의 비교

분석은 대표적인 개선 방식인 Multi-Query Attention (MQA)과 Multi-Head Latent Attention (MLA)에 초점을 맞췄습니다. MQA는 모든 어텐션 헤드가 동일한 키-값 파라미터 세트를 공유하는 보다 적극적인 파라미터 공유 전략을 사용합니다. 이는 메모리 사용량을 줄이지만 모델의 표현력을 저해할 수 있습니다. MLA는 파라미터 압축을 위한 공유 잠재 공간을 도입하지만, 실제 표현력은 가장 작은 차원에 의해 제한되어 중간 차원을 늘려도 성능이 크게 향상되지 않습니다.

MFA의 주요 혁신

MFA 개발은 이론적 성능 한계에 가까워지면서 리소스 소비를 최소화하는 어텐션 메커니즘을 만드는 목표로 추진되었습니다. MFA의 디자인은 세 가지 주요 혁신을 통합합니다.

  • 모델 용량을 극대화하기 위해 어텐션 헤드 수와 차원을 크게 늘립니다.
  • 어텐션 헤드 수와 차원을 확장하면서 파라미터 효율성을 유지하기 위해 적극적인 저랭크 분해 전략을 사용합니다.
  • 모델 복잡성이 증가하더라도 메모리 소비를 최소화하기 위해 단일 키-값 헤드 디자인을 활용합니다.

용량 측정 및 비교

MFA 및 기타 어텐션 메커니즘을 추가로 분석하기 위해 연구팀은 두 가지 주요 메트릭을 도입했습니다.

  • Total Effective Rank (TER): 어텐션 헤드 수와 헤드당 분해 순위(FRH)의 곱입니다.
  • Shared Latent Subspace Dimension (SLSD): 모든 어텐션 헤드가 공유하는 숨겨진 공간의 차원입니다.

MFA는 MQA에 비해 더 높은 SLSD와 TER을 달성합니다. MLA와 비교하여 MFA는 유사한 파라미터 예산으로 더 작은 KV 캐시 크기와 더 높은 TER을 달성하면서도 유사한 SLSD를 유지합니다. 기존 MHA와 비교하여 MFA는 SLSD는 더 작지만 더 높은 TER을 가집니다.

실험 결과

10B에서 1T까지의 학습 데이터와 1B에서 7B 파라미터 범위의 모델을 테스트하여 새로운 아키텍처의 대규모 성능을 평가하기 위한 광범위한 실험을 수행했습니다. MFA는 대규모에서도 뛰어난 성능을 유지하면서 기존 MHA와 유사한 확장 기능을 입증했습니다. MFA-KR은 약간 낮은 성능을 보였지만 확장 추세는 MHA와 일치했습니다. MFA와 MFA-KR의 메모리 절감 이점은 모델 크기가 커짐에 따라 계속 확장되었으며, MFA는 87.5%의 메모리 절감 효과를 달성했고, MFA-KR은 최대 규모에서 메모리 사용량을 6.25%로 줄였습니다.

제거 연구

제거 연구는 MFA 및 MFA-KR의 효율성을 입증했습니다. 성능 이점은 다양한 주류 위치 인코딩 방법에서도 확인되었습니다.

전망

MFA는 단순한 디자인으로 상당한 개선을 제공하여 추가 엔지니어링 복잡성 없이 LLM 추론의 메모리 병목 현상을 효과적으로 해결합니다. 기존 Transformer 생태계에 원활하게 통합되어 다양한 시나리오에서 LLM의 적용을 가속화합니다.