1. 문제 정의 (Problem)

현대 트랜스포머 모델에서 멀티헤드 어텐션 (Multi-Headed Attention, MHA) 은 다음과 같은 두 가지 주요 병목 현상을 겪고 있습니다.

• 계산 복잡도: 어텐션 연산은 시퀀스 길이 ($S$) 와 은닉 차원 ($E$) 에 대해 이차 ($O(S^2)$) 및 선형 ($O(E^2)$) 으로 복잡도가 증가하여 긴 컨텍스트 처리 시 훈련 및 프리필 (prefill) 속도를 저하시킵니다.
• KV 캐시 메모리: autoregressive 생성 시 키 (Key) 와 값 (Value) 캐시가 시퀀스 길이에 비례하여 선형적으로 증가하며, 이는 메모리 대역폭 (Memory Bandwidth) 을 과도하게 소모하여 디코딩 속도를 제한합니다.

기존의 해결책들:

• GQA (Grouped Query Attention): K/V 헤드를 공유하여 KV 캐시 크기를 줄이지만, 훈련 및 프리필 시 필요한 연산량 (FLOPs) 은 MHA 와 동일하게 유지합니다.
• MLA (Multi-Latent Attention): K/V 를 저차원 잠재 공간으로 압축하여 캐시를 줄이지만, 어텐션 연산 전/후에 업 - 프로젝션 (up-projection) 을 수행해야 하므로 계산 비용이 줄어들지 않으며, 오히려 파라미터가 증가할 수 있습니다. 또한 RoPE(위치 임베딩) 처리가 복잡해 별도의 헤드가 필요합니다.

2. 방법론 (Methodology)

저자들은 Compressed Convolutional Attention (CCA) 을 제안하며, 이는 압축된 잠재 공간 (Compressed Latent Space) 에서 전체 어텐션 연산을 수행하는 새로운 접근법입니다.

핵심 기법

1. 완전한 잠재 공간 압축:

   • Query, Key, Value 를 모두 저차원 잠재 공간 ($\tilde{e} = E/C$, 여기서 $C$는 압축 비율) 으로 다운 - 프로젝션합니다.
   • 기존 MLA 와 달리 업 - 프로젝션 없이 압축된 공간에서 직접 어텐션 연산 ($QK^T$ 및 $Attn \cdot V$) 을 수행합니다. 이로 인해 파라미터 수와 FLOPs 가 압축 비율 $C$만큼 감소합니다.
   • RoPE 를 압축된 공간에 직접 적용할 수 있어 별도의 RoPE 헤드가 불필요합니다.

2. 합성곱 혼합 (Convolutional Mixing):

   • 압축된 Query 와 Key 에 대해 시퀀스 (sequence) 및 채널 (channel) 차원에서의 합성곱 (Convolution) 연산을 추가합니다.
   • 이는 잠재 공간에서 학습된 변환에 표현력 (expressivity) 을 부여하고 정보 전달을 원활하게 하여, 단순 압축으로 인한 성능 저하를 보완합니다.

3. QK-Mean 및 Value-Shift:

   • QK-Mean: 컨볼루션 전후의 Q와 K 평균을 더하는 잔차 연결을 통해 정보 공유를 돕고, 어텐션 행렬의 대각선 희소성을 높입니다.
   • Value-Shift: 현재 입력과 이전 입력에서 파생된 두 가지 Value 벡터를 사용하여 헤드의 절반이 현재 토큰을 보지 못하도록 하는 인덕티브 바이어스 (inductive bias) 를 도입합니다. 이는 RNN 기반 시퀀스 믹서 (RWKV 의 Token-shift 등) 에서의 성공적인 사례를 차용한 것입니다.

4. CCGQA (Compressed Convolutional Grouped Query Attention):

   • CCA 에 GQA 의 헤드 공유 전략을 결합한 변형입니다.
   • 잠재 공간 내에서 K/V 헤드를 그룹화하여 공유함으로써, 추가적인 메모리 절감 (최대 2 배) 을 달성하면서도 성능을 유지합니다.
   • Query 와 K/V 의 압축 비율 ($C_1, C_2$) 을 분리하여 유연하게 튜닝할 수 있습니다.

3. 주요 기여 (Key Contributions)

• 계산 및 메모리 동시 최적화: CCA 는 KV 캐시 크기뿐만 아니라 훈련 및 프리필 시 필요한 FLOPs 를 압축 비율에 비례하여 감소시킵니다. 이는 기존 GQA 나 MLA 가 해결하지 못했던 "계산 병목" 문제를 동시에 해결합니다.
• RoPE 통합의 간소화: 압축된 공간에서 직접 RoPE 를 적용하므로, MLA 와 같은 복잡한 별도의 캐시 관리나 헤드가 필요 없습니다.
• Pareto 프론티어 확장: 파라미터 공유 (GQA) 와 파라미터 압축 (MLA) 이 직교 (orthogonal) 한다는 점을 규명하고, 이를 결합한 CCGQA 를 통해 메모리 대역폭 제한과 계산 제한 사이에서 성능을 희생하지 않고 유연하게 트레이드오프할 수 있는 새로운 최적화 곡선을 제시합니다.
• 효율적인 커널 구현: H100 GPU 에서 Flash Attention 스타일의 퓨즈드 커널 (fused kernel) 을 구현하여, 긴 시퀀스 (16k) 에서 MHA 대비 1.7 배 빠른 프리필 및 1.3 배 빠른 백워드 속도를 달성했습니다.

4. 실험 결과 (Results)

• 밀집 모델 (Dense Models):
  • 파라미터 수를 동일하게 맞췄을 때, CCA 는 MLA 보다 낮은 손실 (Loss) 을 기록하며 더 나은 성능을 보였습니다.
  • CCGQA 는 MHA 보다 더 적은 FLOPs 와 KV 캐시로 동등하거나 더 나은 성능을 달성했습니다.
• MoE (Mixture of Experts) 모델:
  • MoE 환경에서 CCGQA 는 GQA 및 MLA 를 능가했습니다.
  • 특히 8 배 KV 캐시 압축을 달성하면서도 MHA 대비 성능 저하가 없었으며, GQA/MLA 대비 반분의 KV 캐시만 사용하면서도 더 높은 성능을 보였습니다.
  • 파라미터를 어텐션 헤드가 아닌 전문가 (Expert) 층으로 재배분할 수 있어 MoE 모델의 효율성을 극대화했습니다.
• 성능 벤치마크:
  • H100 GPU 에서 시퀀스 길이 16k 기준, MHA 대비 프리필 지연 시간 (Latency) 이 약 1.7 배 단축되었습니다.
  • HellaSwag, ARC, Piqa 등 다양한 벤치마크에서 CCGQA 가 MHA, GQA, MLA 보다 일관되게 우수한 결과를 보였습니다.

5. 의의 및 결론 (Significance)

이 논문은 압축된 잠재 공간 전체에서 어텐션을 수행함으로써, 기존 트랜스포머의 계산 복잡도와 메모리 요구 사항을 동시에 획기적으로 줄이는 새로운 패러다임을 제시합니다.

• 긴 컨텍스트 처리: 계산 비용이 줄어들어 동일한 하드웨어 예산으로 더 긴 시퀀스를 처리할 수 있게 되었습니다.
• 유연성: CCA 와 CCGQA 는 훈련 (FLOPs 중점) 과 추론 (메모리 대역폭 중점) 시나리오에 따라 압축 비율을 유연하게 조정할 수 있어, 다양한 하드웨어 제약 조건에 적응 가능합니다.
• 미래 지향성: 시퀀스 압축 기법 (NSA, MoBA 등) 과 결합하여 어텐션 비용을 더욱 낮출 수 있는 가능성을 열어주었으며, 압축된 KV 캐시와 오프라인 압축 기법 간의 상호작용에 대한 후속 연구를 유도합니다.

요약하자면, CCA/CCGQA 는 성능 저하 없이 모델의 크기와 계산 비용을 대폭 줄일 수 있는 차세대 어텐션 메커니즘으로, 특히 긴 컨텍스트와 MoE 아키텍처를 가진 대규모 언어 모델의 효율성을 혁신할 잠재력을 가지고 있습니다.