기술 요약: A3: 어텐션을 위한 분석적 저랭크 근사 프레임워크

문제 정의

대형 언어 모델 (LLM) 은 뛰어난 성능을 달성했으나 방대한 파라미터 수로 인해 높은 계산 비용, 메모리 요구 사항, 그리고 비싼 배포 비용에 시달립니다. 저랭크 근사 ($W \approx A_r B_r$) 는 가중치 행렬을 분해함으로써 유망한 압축 전략을 제시하지만, 기존 접근 방식은 두 가지 중요한 한계에 직면해 있습니다:

1. 아키텍처 무지성: 이들은 어텐션 점수 대 어텐션 출력과 같은 트랜스포머 구성 요소의 특정 기능적 특성을 고려하지 않고 개별 선형 레이어의 출력 오류를 최소화하는 경향이 있습니다.
2. 런타임 오버헤드: 하나의 큰 가중치 행렬을 두 개의 작은 행렬로 분해하여 하나의 일반 행렬 곱셈 (GEMM) 을 두 개의 작은 GEMM 으로 대체합니다. 이는 추가 커널 실행 및 메모리 연산을 포함한 런타임 오버헤드를 도입하여, 종종 가지치기나 양자화보다 열악한 성능을 초래합니다.

결과적으로 현재 저랭크 방법들은 가지치기 및 양자화와 같은 다른 압축 기술의 효율성과 성능을 종종 따라잡지 못합니다.

방법론: A3 프레임워크

저자들은 이러한 한계를 해결하기 위해 트랜스포머 레이어를 세 가지 기능적 구성 요소인 QK(Query-Key), OV(Output-Value), 그리고 MLP(Multi-Layer Perceptron) 로 분해하는 사후 학습 분석적 저랭크 근사 프레임워크인 A3을 제안합니다. 가중치 오류나 일반적인 레이어 출력을 최소화하는 대신, A3 은 각 구성 요소에 고유한 **기능적 손실 (functional loss)**을 최소화하여 숨겨진 차원을 직접 줄이는 폐쇄형 분석적 해를 도출합니다.

1. A3-QK (Query-Key 구성 요소)

• 목표: 소프트맥스 전 어텐션 점수 ($A_i = Q_i K_i^T$) 의 오류를 최소화합니다.
• 방법: 저자들은 캘리브레이션 세트에서 원본과 근사된 어텐션 점수 간의 차이에 대한 푸비니우스 노름 (Frobenius norm) 을 최소화하는 문제로 문제를 공식화합니다.
• 해결책: 그들은 쿼리와 키/값 입력의 자기상관 행렬 ($R_{X_q X_q}$ 및 $R_{X_{kv} X_{kv}}$) 을 포함하는 폐쇄형 해 (정리 3.2) 를 유도합니다. 최적의 저랭크 근사는 융합된 가중치 행렬의 가중 잘린 특이값 분해 (SVD) 를 통해 얻어집니다.
• 확장: 그룹 쿼리 어텐션 (GQA) 의 경우, 동일한 키/값을 공유하는 헤드 간에 결합된 SVD 를 적용합니다. 회전 위치 임베딩 (RoPE) 의 경우, CUR 분해를 사용하여 쌍별 회전을 처리하고 인접한 쌍의 중요도 합으로 차원을 정렬합니다.

2. A3-OV (Output-Value 구성 요소)

• 목표: 어텐션 출력 ($O = \sum A'_i V_i W_o$) 의 오류를 최소화합니다.
• 방법: 높은 계산 비용이 발생하는 글로벌 어텐션 출력 오류를 최소화하는 대신, 프레임워크는 각 헤드의 출력 오류를 독립적으로 최소화하도록 목표를 완화합니다.
• 해결책: 최적 해 (정리 3.3) 는 A3-QK 와 유사하게 유도되지만, 키/값 입력에 곱해진 소프트맥스 후 어텐션 점수 ($P_i$) 의 자기상관에 초점을 맞춥니다. 이는 값과 출력 프로젝션이 공유하는 헤드 차원을 줄입니다.

3. A3-MLP (Multi-Layer Perceptron)

• 목표: MLP 레이어 출력의 오류를 최소화합니다.
• 과제: 비선형 활성화 함수 (예: SiLU, ReLU) 의 존재로 인해 직접적인 SVD 적용이 불가능합니다.
• 방법: 프레임워크는 다운프로젝션 가중치 행렬 ($W_d$) 의 열 부분집합을 선택하여 유지하는 문제로 문제를 공식화하여 중간 차원 ($d_{inter}$) 을 효과적으로 줄입니다.
• 해결책: CUR 분해 원리를 사용하여 저자들은 업프로젝션 가중치의 해당 행의 $L_2$ 노름과 중간 활성화의 자기상관 행렬의 열의 곱을 기반으로 상위 $r$개의 열을 선택합니다. 이는 NP-하드 열 선택에 대한 간단하고 효과적인 휴리스틱입니다.

주요 설계 특징: 런타임 효율성

$(XA)B$로$XW$를 대체하여 추가 GEMM 을 도입하는 전통적인 저랭크 방법과 달리, A3 은 숨겨진 차원 ($d_{qk}$, $d_{vo}$, $d_{inter}$) 을 직접 줄입니다. 이는 원래 모델과 동일한 수의 GEMM 연산을 수행하지만 행렬 크기는 더 작게 만듭니다. 결과적으로 A3 은 추가 커널 실행 없이 모델 크기, KV 캐시 크기, 그리고 FLOPs 를 런타임 오버헤드 없이 줄입니다.

주요 기여

1. 기능적 구성 요소 분해: QK 점수, OV 출력, 그리고 MLP 출력을 위한 일반적인 레이어 출력 대신 기능적 손실 최소화로 최적화를 공식화하는 멀티헤드 어텐션 (MHA) 을 위한 3 부분 저랭크 설정.
2. 분석적 폐쇄형 해: 런타임 오버헤드를 피하면서 자연스럽게 숨겨진 차원을 줄이는 세 가지 목표에 대한 최적 해의 유도.
3. 아키텍처 적응성: 기존 MHA 로 제한된 기존 방법의 한계를 극복하고 그룹 쿼리 어텐션 (GQA) 및 회전 위치 임베딩 (RoPE) 을 포함한 현대적 트랜스포머 변형을 지원하도록 프레임워크 확장.
4. 다용도 응용: KV 캐시 압축, 양자화, 파인튜닝, 그리고 혼합 랭크 할당과 함께 A3 의 호환성 입증.

실험 결과

저자들은 다양한 LLM(LLaMA 1/2/3.1, MPT, Phi 3, Mistral) 에 대해 사전 학습 작업 (WikiText-2, C4, SlimPajama) 과 하류 벤치마크 (ARC, BoolQ, MMLU 등) 에서 광범위한 실험을 수행했습니다.

• 성능 우위: A3 은 최신 (SoTA) 저랭크 방법 (예: SVD-LLM, ASVD, FWSVD) 보다 현저히 우수한 성능을 발휘합니다.
  • 예시: LLaMA 3.1-70B에서 10% 압축 비율로 A3 은 WikiText-2 에서 4.69의 퍼플렉시티를 달성하여, 7.87점을 기록한 이전 SoTA(SVD-LLM) 보다 3.18포인트의 퍼플렉시티 감소를 이루었습니다.
  • LLaMA-2-7B 에서 A3 은 10%, 20%, 40% 압축 비율 전반에 걸쳐 SVD, FWSVD, SVD-LLM 보다 일관되게 더 낮은 퍼플렉시티를 유지합니다.
• 처리량 향상: A3 은 추가 커널 실행 없이 문제 크기를 줄이기 때문에 초당 토큰 수 (TPS) 에서 일관된 속도 향상을 달성합니다. 반면, SVD-LLM 과 같은 방법들은 종종 속도 패널티를 겪거나 공격적인 압축 수준에서만 이점을 보입니다.
• KV 캐시 효율성: A3 은 Palu 및 CLOVER 와 같은 방법에 비해 더 강력한 KV 캐시 압축을 제공하여, 높은 압축 비율 (예: MPT-30B 에서 80%) 에서도 더 낮은 퍼플렉시티를 유지합니다.
• 호환성: A3 은 양자화 (예: HQQ) 와 효과적으로 결합하여 압축 수준의 연속 스펙트럼을 생성하고, 4 비트 미만 영역에서 양자화 단독보다 더 나은 파레토 프런티어를 달성합니다.

중요성과 주장

이 논문은 A3 이 기능적 구성 요소 오류로 최적화 초점을 이동시킴으로써 모델 압축 분야에서 중요한 진전을 이루었다고 주장합니다. 숨겨진 차원을 직접 줄이는 분석적 해를 도출함으로써 A3 은 다음을 달성합니다:

• 우수한 정확도: 기존 저랭크 근사 방법보다 훨씬 더 나은 엔드투엔드 모델 성능 유지.
• 하드웨어 효율성: 분해된 행렬 곱셈과 관련된 런타임 패널티 없이 모델 크기, 메모리 풋프린트 (KV 캐시 포함), 그리고 FLOPs 감소.
• 광범위한 적용 가능성: 이전 저랭크 방법들이 효과적으로 처리하지 못했던 현대 아키텍처 (GQA, RoPE) 에 성공적으로 적응.

저자들은 A3 이 런타임 오버헤드를 최소화해야 하는 시나리오에서 가지치기 및 양자화에 대한 실용적이고 고성능의 대안을 제공하는 사후 학습, 학습 없는 방법임을 강조합니다.