[논문 리뷰] GQA 기반으로 학습된 LLM에 Multi-Head Latent Attention 적용해보기

 

이전 게시물에서 간단하게 Multi-Head Latent Attention에 대해 알아보았으며, 본 게시물에서는 "TransMLA: Multi-Head Latent Attention Is All You Need" 이라는 논문을 통해 Grouped Query Attention (GQA)를 기반으로 학습된 대형 언어 모델에 MLA를 어떻게 적용할 수 있는지에 대해서 알아보고자 합니다. Multi-Head Latention Attention에 대해 이해가 안되신다면, 아래의 게시물 참고 부탁드립니다. 

"Attention Is All You Need" 의 대항마 : Multi-Head Latent Attention

 

본 게시물에 대해 직관적으로 이해하기 위해서는 Singular Value Decompostion에 대한 개념이 있어야 하므로, 이에 대해 궁금하시다면 숨긴 글 참조 부탁드립니다. 

Singular Value Decomposition (SVD)는 행렬 대수학에서 매우 유용하게 사용되는 행렬 분해 방법 중 하나이다. 임의의 m×n 크기의 행렬 A를 세 개의 특별한 행렬의 곱, 즉 A=UΣVT 형태로 분해하게 된다. 여기서 U는 m×m 크기의 직교 행렬(열벡터들이 서로 직교), V는 n×n 크기의 직교 행렬(열벡터들이 서로 직교), 그리고 Σ는 m×n 크기의 대각 행렬로, 대각선 상의 값들을 singular values이라고 칭한다. 이 특이값들은 해당 방향으로 데이터가 얼마나 큰 분산을 가지는지, 즉 얼마나 중요한지를 의미한다. 머신러닝에 대해 다뤄본 경험이 있다면, 한번쯤은 들어보았을 만한 PCA 등과 굉장히 비슷한 방법으로 사용이 된다. SVD는 차원 축소, 노이즈 제거, 데이터의 주요 특징 추출 등 다양한 분야에 활용되며, TransMLA에서는 프로젝션 행렬들의 직교성을 확보하고 중요한 정보를 유지하면서도 중복성을 줄이는 데  가장 핵심적인 역할을 한다. 

---

GQA에서 MLA의 변환 과정: 어떻게 가능할까?

이미지 출처: TransMLA 논문

Multi-Head Latent Attention(MLA)는 Grouped Query Attention(GQA)**에 비해 훨씬 높은 표현력을 갖는다. 여기서 표현력이란, 모델이 다양한 패턴과 관계를 학습하고 이를 유연하게 표현할 수 있는 능력을 의미한다. 직관적으로 GQA는 query head의 수는 상대적으로 많지만, key-value head의 수는 적기 때문에, 다수의 query가 동일한 key-value 표현을 공유하게 된다. 이러한 구조는 계산 효율성 측면에서는 이점을 가지나, 표현의 다양성과 유연성 측면에서는 한계를 지닌다.

 

반면, MLA는 Multi-Head Attention(MHA)처럼 각 attention head가 독립적인 key-value 행렬을 가지도록 설계되어, head 간 상호작용이 자유롭고 보다 복잡한 관계를 모델링할 수 있다. 이러한 설계는 attention의 표현 공간을 확장하며, downstream task에서의 성능 향상으로 이어질 수 있다.

 

본 논문에서는 GQA의 구조적 틀은 유지하되, 내부 계산 방식—특히 key와 value의 생성 및 저장 구조—를 MLA 방식으로 재설계한 TransMLA 구조를 제안한다. 기존 GQA는 적은 수의 key-value head를 계산한 후 이를 query head 수에 맞춰 단순 복제(Replication)함으로써 계산량을 줄이는 방식이다. 하지만 이러한 복제 과정은 서로 다른 query들이 동일한 key-value 표현을 공유하게 만들며, 이는 모델의 표현력 측면에서 주요한 병목으로 작용한다.

 

TransMLA는 GQA의 복제된 구조를 제거하고, key-value를 생성하는 가중치 행렬을 저차원(latent) 공간에서 factorization하는 방식으로 설계되었다. 이를 통해 각 query head는 서로 다른 key-value 표현을 참조할 수 있게 되며, MLA 수준의 head-wise 다양성과 표현력을 확보할 수 있게 되었다. 이러한 구조적 전환은 GQA의 계산 효율성을 유지하면서도, MLA의 핵심 장점인 다양한 표현 학습 능력과 복잡한 관계 모델링 능력을 함께 실현할 수 있다. 결과적으로, TransMLA는 기존 GQA 기반의 LLM 구조에 최소한의 수정만으로도 성능을 향상시킬 수 있는 새로운 attention 설계 방식으로 제안하였다. 

---

GQA에서 MLA 변환 과정: TransMLA 방법론 

 

TranMLA는 트랜스포머 모델의 GQA을 키-값 헤드 수를 최적화하여, MLA 형태로 변화하는 방법론을 제안한다. 이러한 변환의 핵심은 키-값 프로젝션에 직교성을 부여하여 SVD의 수학적 기법을 활용함으로써, 동일한 KV 캐시 크기 내에서 모델의 표현력을 향샹시키는 것이 가장 핵심적이다.  

 

어텐션 메커니즘의 핵심 구성 요소인 쿼리(Q), 키(K), 값(V) 프로젝션은 학습 가능한 가중치 행렬들로, 입력 토큰 시퀀스를 각각의 역할에 맞는 고차원 벡터 공간으로 변환하게 된다. 예를 들어, 특정 단어의 Q 벡터는 "질문"을 던지는 역할을 하며, 다른 단어들의 K 벡터들과의 유사도를 계산하여 어떤 정보에 집중할지를 결정한다. 이때 Q,K,V 프로젝션들이 서로 직교성을 갖도록, 즉 각 프로젝션이 만드는 벡터 공간의 주요 축들이 서로 90도에 가깝게 유지되도록 하는 것은 매우 중요하다. 조금 더 쉽게 말하자면, 두 벡터 간의 내적이 0에 가깝다는 것을 의미하고, 이는 두 정보 간의 상관관계가 매우 낮거나 독립적이라고 표현할 수 있다.

결국, Q,K,V 프로젝션, 특히 K,V 프로젝션에서 직교성을 강제하는 것은 제한된 계산 자원 내에서 정보의 중복을 최소화하고, 각 어텐션 헤드가 가능한 한 독립적이고 유의미한 정보를 처리함으로써, 모델 전체의 표현력, 학습 안정성 등과 같은 성능적인 측면에서 향상시키기 위한 근본적인 전략이라고 생각하는 것이 좋을거 같다. 

 

위와 같은 내용은 TransMLA 논문에서 대체적으로 비슷하게 적용이 된다. SVD는 주어진 행렬 A를 UΣV^T로 분해하는데, 여기에서 U와 V의 직교 기저 벡터들을 활용하여 K와 V 프로젝션 가중치를 재구성하거나, 혹은 이들 통해 새로운 프로젝션 레이어를 형성하게 된다. 이런 방향으로 진행함으로써, K와 V 벡터들이 생성되는 공간의 주요 방향들이 서로 최대한 독립적이 되도록 강제하여, 서로 다른 특징을 포작하더록 강제한다. 이를 통해 만약 프로젝션들이 서로 강하게 상관되어 있다면, 특정 파라미터의 작은 변화가 다른 많은 부분에 연쇄적인 큰 영향을 미쳐 최적화 과정에서 그래디언트가 폭주하거나 소실될 위험이 커지지만,  직교성 강제하는 것은 상호 의존성을 줄여 최적화 과정을 원활하고 예측 가능하게 만들어 준다. 하지만 반대로 프로젝션들이 서로 유사한 방향으로 치우쳐져 있다면 (구체적으로 직교성이 낮다면), 여러 어텐션 헤드가 실제로는 거의 같은 정보에 집중하여 중복성이 발생할 수 있다. 이는 서로 유사한 정보를 중복적으로 처리하게 되어 어텐션 헤드의 중복성을 야기할 수 있다.

 

또한 MLA와 같이 쿼리 헤드 수에 비해 키-값 헤드 수가 상대적으로 작은 어텐션 구조에서는 이러한 직교성의 확보가 중요해진다. 제한된 수의 키-값 헤드가 다양한 쿼리에 대해 충분한 표현력을 제공하지 못하게 된다면, 서로 다른 쿼리 입력에 대해 유사한 결과만을 도출하는 Representational Collapse 라는 문제점이 발생할 가능성이 매우 커지게 된다. 이에 TransMLA는 SVD를 통해 키와 값 프로젝션을 정제함으로써, 적은 수의 어텐션 헤드만으로도 정보의 다양성을 유지하고 메모리 효율성을 달상하며, 동시에 데이터 내의 복잡한 관계까지 효과적으로 포착할 수 있는 유망한 어텐션 기법이다.