[Paper] 멀티헤드 Low-Rank 어텐션

Source: arXiv - 2603.02188v1

개요

The paper introduces Multi-Head Low-Rank Attention (MLRA), a new attention mechanism designed to speed up long‑context decoding in large language models (LLMs). By making the latent representation partitionable across multiple GPUs, MLRA cuts the memory traffic that normally bottlenecks token‑by‑token generation, delivering up to a 2.8× speed boost while preserving (or even improving) model quality.

주요 기여

• 분할 가능한 잠재 상태: Multi‑Head Latent Attention (MLA)와 달리, MLRA의 저‑랭크 잠재 벡터는 디바이스 간에 분할될 수 있어 효율적인 Tensor Parallel (TP) 디코딩이 가능하다.
• 4‑way TP 친화적 설계: 아키텍처는 각 GPU가 KV 캐시의 자신에게 할당된 조각만 로드하도록 하여 오프‑칩 메모리 대역폭 사용을 크게 감소시킨다.
• 최첨단 성능: 실험 결과 MLRA가 퍼플렉시티와 QA, 요약 등 다운스트림 벤치마크에서 MLA와 동등하거나 능가함을 보여준다.
• 디코딩 속도 향상: 동일한 하드웨어 구성에서 MLA 대비 2.8× 디코딩 가속을 측정했다.
• 오픈‑소스 공개: 코드, 사전 학습된 가중치, 학습/평가 스크립트가 공개되어 재현성과 커뮤니티 채택을 촉진한다.

방법론

1. Low‑rank factorisation of attention

   • 전통적인 자체 주의는 각 토큰마다 전체 키‑값(KV) 행렬을 저장하며, 이는 매 생성 단계마다 고대역폭 메모리(HBM)에서 가져와야 합니다.
   • MLRA는 KV 캐시를 잠재 저차원 표현(rank‑r)과 프로젝션 행렬 집합으로 분해하여 전체 캐시 크기를 감소시킵니다.

2. Multi‑head design with sharding support

   • 이제 각 어텐션 헤드는 자체 저랭크 잠재 상태를 포함합니다. 잠재 차원이 헤드들 사이에 분할되므로, 이 상태를 TP 설정에서 GPU들 간에 샤딩할 수 있습니다.
   • 디코딩 중에는 각 디바이스가 자신이 소유한 잠재 캐시 조각만 로드하고, 프로젝션 행렬은 칩 내에 유지됩니다.

3. Training pipeline

   • 저자들은 MLRA 모듈을 삽입한 표준 디코더 전용 트랜스포머(예: LLaMA와 유사)를 사전 학습합니다.
   • 인과 언어 모델링 손실과 저랭크 요인이 대부분의 어텐션 정보를 포착하도록 유도하는 정규화 항을 결합하여 사용합니다.

4. Evaluation

   • 벤치마크에는 표준 코퍼스에 대한 언어 모델링 퍼플렉시티와 MMLU, GSM‑8K, 요약 스위트와 같은 다운스트림 작업이 포함됩니다.
   • 디코딩 속도는 4‑GPU 노드(NVIDIA A100, 80 GB)에서 탐욕적 검색과 빔 서치를 모두 사용하여 측정합니다.

결과 및 발견

• 품질: MLRA는 베이스라인 및 MLA 모두에 비해 퍼플렉시티와 다운스트림 정확도가 약간 향상되어 저‑랭크 팩터라이제이션이 표현력을 희생하지 않음을 확인합니다.
• 속도: 4‑way TP 구현은 KV 캐시 트래픽을 약 65 % 감소시켜 MLA 대비 2.8× 디코딩 속도 향상을, full‑KV 베이스라인 대비 ~2× 향상을 제공합니다.
• 확장성: 2 GPU에서 8 GPU까지 확장한 실험에서 거의 선형적인 처리량 증가를 보여주며, 샤딩 설계가 의도대로 작동함을 입증합니다.

Practical Implications

• Faster inference for long‑context apps: 수천 개의 토큰을 컨텍스트에 유지해야 하는 챗봇, 코드 어시스턴트, 문서 분석 도구가 이제 응답을 눈에 띄게 빠르게 생성할 수 있어 최종 사용자의 지연 시간이 감소합니다.
• Lower hardware cost per token: 각 GPU가 KV 캐시의 일부만 가져오면 되므로, 동일한 추론 처리량을 더 저렴한 GPU 클러스터에서 혹은 오프‑칩 트래픽이 감소된 단일 고메모리 GPU에서도 달성할 수 있습니다.
• Improved TP utilisation: 기존 텐서‑패럴렐 파이프라인(예: DeepSpeed, Megatron‑LM)이 최소한의 변경만으로 MLRA를 도입할 수 있어, 메모리 효율성과 속도를 동시에 얻으면서 가중치 샤딩과 같은 모델 패럴렐리즘 이점을 유지합니다.
• Easier deployment of LLMs on edge‑like servers: 메모리 대역폭 요구가 감소함에 따라 HBM이 제한된 서버에서도 대형 모델을 실행할 수 있게 되어, 온‑프레미스 또는 프라이빗 클라우드 LLM 서비스 제공이 가능해집니다.

Limitations & Future Work

• Rank selection sensitivity: 낮은 차원의 랭크(r)는 신중하게 조정해야 합니다; 너무 낮으면 정확도가 떨어지고, 너무 높으면 메모리 트래픽 이득이 감소합니다. 자동 랭크 선택 전략은 탐구되지 않았습니다.
• Focus on decoder‑only transformers: 이 논문은 인과 언어 모델만을 평가합니다. MLRA를 인코더‑디코더 아키텍처(예: T5)나 비전‑언어 모델에 확장하는 것은 아직 미해결 과제입니다.
• Hardware‑specific optimisations: 보고된 속도 향상은 A100 GPU에서 측정되었습니다; 다른 가속기(예: AMD GPU, TPU)에서의 성능은 다를 수 있으며 전용 커널 튜닝이 필요합니다.
• Training overhead: 저랭크 팩터화를 도입하면 추가 투영 레이어가 생겨 학습 계산량이 약간 증가합니다. 향후 작업에서는 보다 효율적인 학습 트릭이나 혼합 정밀도 방식을 조사할 수 있습니다.

전반적으로, Multi‑Head Low‑Rank Attention은 모델 품질을 희생하지 않으면서 고처리량, 장기 컨텍스트 LLM 추론이 필요한 개발자에게 매력적인 레시피를 제공합니다. 코드와 사전 학습된 가중치가 이미 공개되어 있어, 커뮤니티는 바로 실험을 시작할 수 있습니다.

저자

• Songtao Liu
• Hongwu Peng
• Zhiwei Zhang
• Zhengyu Chen
• Yue Guo

논문 정보

• arXiv ID: 2603.02188v1
• Categories: cs.LG
• Published: 2026년 3월 2일
• PDF: PDF 다운로드